import numpy as np
from sklearn.utils.multiclass import type_of_target
from mindware.base_estimator import BaseEstimator
from mindware.components.utils.constants import type_dict, MULTILABEL_CLS, IMG_CLS, TEXT_CLS, OBJECT_DET
from mindware.components.feature_engineering.transformation_graph import DataNode


class Classifier(BaseEstimator):
    """This class implements the classification task. """

    def initialize(self, data: DataNode, **kwargs):
        if self.metric is None:
            self.metric = 'acc'

        # Check the task type: {binary, multiclass}
        task_type = type_of_target(data.data[1])
        if task_type in type_dict:
            task_type = type_dict[task_type]
        else:
            raise ValueError("Invalid Task Type: %s!" % task_type)
        self.task_type = task_type
        super().initialize(data=data, **kwargs)

    def fit(self, data: DataNode, **kwargs):
        """
        Fit the classifier to given training data.
        :param data: instance of DataNode
        :return: self
        """
        if self._ml_engine is None:
            self.initialize(data=data, **kwargs)
        super().fit(data, **kwargs)

        return self

    def predict(self, X, batch_size=None, n_jobs=1):
        """
        Predict classes for X.
        :param X: Datanode
        :param batch_size: int
        :param n_jobs: int
        :return: y : array of shape = [n_samples]
            The predicted classes.
        """
        if not isinstance(X, DataNode):
            raise ValueError("X is supposed to be a Data Node, but get %s" % type(X))
        return super().predict(X, batch_size=batch_size, n_jobs=n_jobs)

    def refit(self):
        return super().refit()

    def predict_proba(self, X, batch_size=None, n_jobs=1):
        """
        Predict probabilities of classes for all samples X.
        :param X: Datanode
        :param batch_size: int
        :param n_jobs: int
        :return: y : array of shape = [n_samples, n_classes]
            The predicted class probabilities.
        """
        if not isinstance(X, DataNode):
            raise ValueError("X is supposed to be a Data Node, but get %s" % type(X))
        pred_proba = super().predict_proba(X, batch_size=batch_size, n_jobs=n_jobs)

        if self.task_type != MULTILABEL_CLS:
            assert (
                np.allclose(
                    np.sum(pred_proba, axis=1),
                    np.ones_like(pred_proba[:, 0]))
            ), "Prediction probability does not sum up to 1!"

        # Check that all probability values lie between 0 and 1.
        assert (
                (pred_proba >= 0).all() and (pred_proba <= 1).all()
        ), "Found prediction probability value outside of [0, 1]!"

        return pred_proba

    def get_tree_importance(self, data: DataNode):
        from lightgbm import LGBMClassifier
        import pandas as pd
        X, y = self.data_transformer(data).data
        lgb = LGBMClassifier(random_state=1)
        lgb.fit(X, y)
        _importance = lgb.feature_importances_
        h = {}
        h['feature_id'] = np.array(range(len(_importance)))
        h['feature_importance'] = _importance
        return pd.DataFrame(h)

    def get_linear_importance(self, data: DataNode):
        from sklearn.linear_model import LogisticRegression
        import pandas as pd
        X, y = self.data_transformer(data).data
        clf = LogisticRegression(random_state=1)
        clf.fit(X, y)
        _ef = clf.coef_
        std_array = np.std(_ef, ddof=1, axis=0)
        abs_array = abs(_ef)
        mean_array = np.mean(abs_array, axis=0)
        _importance = std_array / mean_array
        h = {}
        h['feature_id'] = np.array(range(len(_importance)))
        h['feature_importance'] = _importance
        return pd.DataFrame(h)

    def get_linear_impact(self, data: DataNode):
        from sklearn.linear_model import LogisticRegression
        import pandas as pd
        if (len(set(data.data[1]))) > 2:
            print('ERROR! Only binary classification is supported!')
            return 0
        X, y = self.data_transformer(data).data
        clf = LogisticRegression(random_state=1)
        clf.fit(X, y)
        _ef = clf.coef_
        _impact = _ef[0]
        h = {}
        h['feature_id'] = np.array(range(len(_impact)))
        h['feature_impact'] = _impact
        return pd.DataFrame(h)


class Regressor(BaseEstimator):
    """This class implements the regression task. """

    def initialize(self, data: DataNode, **kwargs):
        self.metric = 'mse' if self.metric is None else self.metric

        # Check the task type: {continuous}
        task_type = type_dict['continuous']
        self.task_type = task_type
        super().initialize(data=data, **kwargs)

    def fit(self, data, **kwargs):
        """
        Fit the regressor to given training data.
        :param data: DataNode
        :return: self
        """
        if self._ml_engine is None:
            self.initialize(data=data, **kwargs)
        super().fit(data, **kwargs)

        return self

    def predict(self, X, batch_size=None, n_jobs=1):
        """
        Make predictions for X.
        :param X: DataNode
        :param batch_size: int
        :param n_jobs: int
        :return: y : array of shape = [n_samples] or [n_samples, n_labels]
            The predicted classes.
        """
        if not isinstance(X, DataNode):
            raise ValueError("X is supposed to be a Data Node, but get %s" % type(X))
        return super().predict(X, batch_size=batch_size, n_jobs=n_jobs)

    def get_tree_importance(self, data: DataNode):
        from lightgbm import LGBMRegressor
        import pandas as pd
        X, y = self.data_transformer(data).data
        lgb = LGBMRegressor(random_state=1)
        lgb.fit(X, y)
        _importance = lgb.feature_importances_
        h = {}
        h['feature_id'] = np.array(range(len(_importance)))
        h['feature_importance'] = _importance
        return pd.DataFrame(h)

    def get_linear_impact(self, data: DataNode):
        from sklearn.linear_model import LinearRegression
        import pandas as pd
        X, y = self.data_transformer(data).data
        reg = LinearRegression()
        reg.fit(X, y)
        _impact = reg.coef_
        h = {}
        h['feature_id'] = np.array(range(len(_impact)))
        h['feature_impact'] = _impact
        return pd.DataFrame(h)

[((16, 20, 16, 48), 'sklearn.utils.multiclass.type_of_target', 'type_of_target', ({(16, 35, 16, 47): 'data.data[1]'}, {}), '(data.data[1])', False, 'from sklearn.utils.multiclass import type_of_target\n'), ((83, 14, 83, 44), 'lightgbm.LGBMClassifier', 'LGBMClassifier', (), '', False, 'from lightgbm import LGBMClassifier\n'), ((89, 15, 89, 30), 'pandas.DataFrame', 'pd.DataFrame', ({(89, 28, 89, 29): 'h'}, {}), '(h)', True, 'import pandas as pd\n'), ((95, 14, 95, 48), 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', (), '', False, 'from sklearn.linear_model import LogisticRegression\n'), ((98, 20, 98, 47), 'numpy.std', 'np.std', (), '', True, 'import numpy as np\n'), ((100, 21, 100, 47), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((105, 15, 105, 30), 'pandas.DataFrame', 'pd.DataFrame', ({(105, 28, 105, 29): 'h'}, {}), '(h)', True, 'import pandas as pd\n'), ((114, 14, 114, 48), 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', (), '', False, 'from sklearn.linear_model import LogisticRegression\n'), ((121, 15, 121, 30), 'pandas.DataFrame', 'pd.DataFrame', ({(121, 28, 121, 29): 'h'}, {}), '(h)', True, 'import pandas as pd\n'), ((164, 14, 164, 43), 'lightgbm.LGBMRegressor', 'LGBMRegressor', (), '', False, 'from lightgbm import LGBMRegressor\n'), ((170, 15, 170, 30), 'pandas.DataFrame', 'pd.DataFrame', ({(170, 28, 170, 29): 'h'}, {}), '(h)', True, 'import pandas as pd\n'), ((176, 14, 176, 32), 'sklearn.linear_model.LinearRegression', 'LinearRegression', ({}, {}), '()', False, 'from sklearn.linear_model import LinearRegression\n'), ((182, 15, 182, 30), 'pandas.DataFrame', 'pd.DataFrame', ({(182, 28, 182, 29): 'h'}, {}), '(h)', True, 'import pandas as pd\n'), ((68, 20, 68, 46), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((69, 20, 69, 50), 'numpy.ones_like', 'np.ones_like', ({(69, 33, 69, 49): 'pred_proba[:, (0)]'}, {}), '(pred_proba[:, (0)])', True, 'import numpy as np\n')]