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)
        num_colors = nclusters
        colors = ['#' + ''.join([random.choice('0123456789ABCDEF') for _ in range(6)]) for _ in range(num_colors)]
        col = np.array(['#' + ''.join([random.choice('0123456789ABCDEF') for _ in range(6)]) for _ in range(self.x.shape[0])])
        for i in range(nclusters):
            ss = np.where(yp==i)
            col[ss] = colors[i]


        return self.x, col