from Packages import *

class Sk_Kmeans:
    def __init__(self, x, max_clusters):
        self.x = x
        self.max_clusters = max_clusters

        self.inertia = pd.DataFrame()
        for i in range(1, max_clusters+1):
            model = KMeans(n_clusters = i, init = 'k-means++', random_state = 42)
            model.fit(x)
            self.inertia[f'{i}_clust']= [model.inertia_]
        self.inertia.index = ['inertia']

    @property
    def inertia_(self):
        return self.inertia

    def fit_optimal(self, nclusters):
        model = KMeans(n_clusters = nclusters, init = 'k-means++', random_state = 42)
        model.fit(self.x)
        yp = model.predict(self.x)+1
        clu = [f'cluster#{i}' for i in yp]
        return self.x, clu