
from sklearn.cluster import KMeans
import numpy as np
from lmi_proteins.classifiers.DummyClassifier import LMIDummyClassifier


def train_kmeans_classifier(classifier, params, X, node_id):
    if X.shape[0] <= params['n_clusters']:
        params['n_clusters'] = classifier.decrease_n_components(
            X.shape[0]
        )

    if params['n_clusters'] <= 1:
        classifier = LMIDummyClassifier(node_id)
        y = [0 for _ in range(X.shape[0])]
        return classifier.fit(X, y)
    else:
        return classifier.fit(X)


class LMIKMeans(KMeans):

    def __init__(self, node_id, **kwargs):
        super().__init__(**kwargs)
        self.node_id = node_id
        self.classes_ = [component for component in range(kwargs['n_clusters'])]

    def decrease_n_components(self, n_objects):
        decreased_value = n_objects // 2
        self.n_clusters = decreased_value
        self.classes_ = [component for component in range(self.n_clusters)]
        return decreased_value

    def fit(
        self,
        X: np.array
    ):
        """
        Trains a K-Means classifier.

        Parameters
        -------
        X: np.array (2D)
            Training values

        Returns
        -------
        predictions : np.array
            Array of model predictions
        """
        super().fit(X)
        self.node_classes_ = [self.node_id + (node, ) for node in self.classes_]
        return super().predict(X)

    def predict(self, query):
        """ Collects predictions for a query (= one object/data point).

        Parameters
        -------
        query: np.array
            - one dimension along x-axis, array looks like: [[descriptors...]]
            Query - expected shape: (1, n_descriptors)
        level_mapping: List
            List of classes in the dataset

        Returns
        -------
        np.array
            Array of model predictions with corresponding classes.
        """
        assert query.shape[0] == 1
        dist_distr = self.transform(query)[0]
        dist_distr = max(dist_distr) - dist_distr
        dist_distr += 1e-8
        dist_distr /= sum(dist_distr)
        assert len(self.node_classes_) == dist_distr.shape[0]
        return np.array((self.node_classes_, dist_distr), dtype=object).T
