import numpy as np
from sklearn.linear_model import LogisticRegression
import logging

from warnings import simplefilter
from sklearn.exceptions import ConvergenceWarning
simplefilter("ignore", category=ConvergenceWarning)
# Filter out convergence warning:
#   Since running LogReg on lots of data and lot of iterations
#   may be expensive, we allow the failure of convergence
#   and needn't notify about it in the logs.

LOG = logging.getLogger(__name__)
LOG.setLevel('DEBUG')


class LMILogisticRegression(LogisticRegression):

    def __init__(
        self,
        node_id,
        max_iter,
        class_weight='balanced',
        solver='lbfgs',
        multi_class='ovr',
        **kwargs
    ):
        super().__init__(
            max_iter=max_iter,
            class_weight=class_weight,
            solver=solver,
            multi_class=multi_class,
            **kwargs
        )
        self.node_id = node_id

    def fit(
        self,
        X: np.array,
        y: np.array
    ):
        """ Trains a Logistic regression model.
        Expects model_dict to contain hyperparameter *ep* (number of epochs)

        Parameters
        -------
        X: Numpy array
            Training values
        y: Numpy array
            Training labels

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

    def predict(self, query):
        assert query.shape[0] == 1
        predictions = self.predict_proba(query)[0]
        assert len(self.node_classes_) == predictions.shape[0]
        return np.array((self.node_classes_, predictions), dtype=object).T
