
from sklearn.mixture import GaussianMixture
import numpy as np
from lmi_proteins.classifiers.DummyClassifier import LMIDummyClassifier


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

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


class LMIGaussianMixtureModel(GaussianMixture):

    def __init__(self, node_id, **kwargs):
        super().__init__(**kwargs)
        self.node_id = node_id
        self.classes_ = [component for component in range(kwargs['n_components'])]
    """
    GMM classifier representing an LMI node.
    """
    def decrease_n_components(self, n_objects):
        decreased_value = n_objects // 2
        self.n_components = decreased_value
        self.classes_ = [component for component in range(self.n_components)]
        return decreased_value

    def fit(
        self,
        X: np.array
    ):
        """
        Creates and trains a RandomForest classifier.
        Expects model_config to contain hyperparameters *depth* and *n_estimators*
            == depth and number of estimators (trees)

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

        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
        predictions = super().predict_proba(query)[0]
        assert len(self.node_classes_) == predictions.shape[0]
        return np.array((self.node_classes_, predictions), dtype=object).T
