# Copyright 2023-2024 The SapientML Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import collections
import os
import pickle
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from sapientml.params import Task

from ..design import search_space
from ..params import DatasetSummary, PipelineSkeleton

name_to_label_mapping = {
    "random forest": {
        "c": "MODEL:Classifier:RandomForestClassifier:sklearn",
        "r": "MODEL:Regressor:RandomForestRegressor:sklearn",
    },
    "extra tree": {
        "c": "MODEL:Classifier:ExtraTreesClassifier:sklearn",
        "r": "MODEL:Regressor:ExtraTreesRegressor:sklearn",
    },
    "lightgbm": {"c": "MODEL:Classifier:LGBMClassifier:lightgbm", "r": "MODEL:Regressor:LGBMRegressor:lightgbm"},
    "xgboost": {"c": "MODEL:Classifier:XGBClassifier:xgboost", "r": "MODEL:Regressor:XGBRegressor:xgboost"},
    "catboost": {
        "c": "MODEL:Classifier:CatBoostClassifier:catboost",
        "r": "MODEL:Regressor:CatBoostRegressor:catboost",
    },
    "gradient boosting": {
        "c": "MODEL:Classifier:GradientBoostingClassifier:sklearn",
        "r": "MODEL:Regressor:GradientBoostingRegressor:sklearn",
    },
    "adaboost": {"c": "MODEL:Classifier:AdaBoostClassifier:sklearn", "r": "MODEL:Regressor:AdaBoostRegressor:sklearn"},
    "decision tree": {
        "c": "MODEL:Classifier:DecisionTreeClassifier:sklearn",
        "r": "MODEL:Regressor:DecisionTreeRegressor:sklearn",
    },
    "svm": {"c": "MODEL:Classifier:SVC:sklearn", "r": "MODEL:Regressor:SVR:sklearn"},
    "linear svm": {"c": "MODEL:Classifier:LinearSVC:sklearn", "r": "MODEL:Regressor:LinearSVR:sklearn"},
    "logistic/linear regression": {
        "c": "MODEL:Classifier:LogisticRegression:sklearn",
        "r": "MODEL:Regressor:LinearRegression:sklearn",
    },
    "lasso": {"r": "MODEL:Regressor:Lasso:sklearn"},
    "sgd": {"c": "MODEL:Classifier:SGDClassifier:sklearn", "r": "MODEL:Regressor:SGDRegressor:sklearn"},
    "mlp": {"c": "MODEL:Classifier:MLPClassifier:sklearn", "r": "MODEL:Regressor:MLPRegressor:sklearn"},
    "multinomial nb": {"c": "MODEL:Classifier:MultinomialNB:sklearn"},
    "gaussian nb": {"c": "MODEL:Classifier:GaussianNB:sklearn"},
    "bernoulli nb": {"c": "MODEL:Classifier:BernoulliNB:sklearn"},
}


class Predicate:
    """Return Parameters in one sentence.

    Parameters
    ----------
    feature_name : str
        Feature_name of the input dataframe.
    value : str
        Value of the input dataframe.
    threshold : str
        Threshold of the input dataframe.
    op : str
        Threshold sign. Evaluation of the value against the threshold.

    Returns
    ----------
    str

    """

    def __init__(self, feature_name, value, threshold, op):
        self.feature_name = feature_name
        self.value = value
        self.threshold = threshold
        self.op = op

    def __repr__(self):
        return self.feature_name + "(" + self.value + ")" + " " + self.op + " " + self.threshold

    def __str__(self):
        return self.feature_name + "(" + self.value + ")" + " " + self.op + " " + self.threshold


def get_decision_path(clf, X):
    """Return a threshold evaluation of the feature value.

    Parameters
    ----------
    X : pd.DataFrame
        Input dataFrame.

    Returns
    ----------
    rules : list
        Threshold evaluation of feature values.
    """
    feature = clf.tree_.feature
    threshold = clf.tree_.threshold

    node_indicator = clf.decision_path(X)
    leaf_id = clf.apply(X)

    rules = []

    for row_index in range(X.shape[0]):
        sample_id = row_index
        # obtain ids of the nodes `sample_id` goes through, i.e., row `sample_id`
        node_index = node_indicator.indices[node_indicator.indptr[sample_id] : node_indicator.indptr[sample_id + 1]]

        predicates = []
        for node_id in node_index:
            # continue to the next node if it is a leaf node
            if leaf_id[sample_id] == node_id:
                continue

            # check if value of the split feature for sample 0 is below threshold
            if X.iloc[sample_id, feature[node_id]] <= threshold[node_id]:
                threshold_sign = "<="
            else:
                threshold_sign = ">"

            predicates.append(
                Predicate(
                    X.columns[feature[node_id]], X.iloc[sample_id, feature[node_id]], threshold[node_id], threshold_sign
                )
            )
        rules.append(predicates)
    return rules


def _predict_preprocessors(pp_models, meta_features: pd.DataFrame, target_labels: list[str]) -> PipelineSkeleton:
    for col in meta_features.columns:
        meta_features[col].fillna(0, inplace=True)
    output = meta_features[search_space.meta_feature_list].copy()
    labels: list[str] = []
    rules = {}
    for label, (model, relevant_features) in pp_models.items():
        if label not in target_labels:
            continue
        labels.append(label)

        final_meta_features = meta_features[relevant_features]

        prediction = model.predict(final_meta_features)

        prediction_proba = model.predict_proba(final_meta_features)[:, 1]
        rules[label] = get_decision_path(model, final_meta_features)
        output[label] = prediction
        output[label + "_proba"] = prediction_proba

    pp_labels = collections.OrderedDict()
    for index, row in output.iterrows():
        labels_info = {}
        for label in labels:
            predicates_obj = rules[label][index]
            predicates = []
            for predicate_obj in predicates_obj:
                predicates.append(
                    {
                        "feature_name": predicate_obj.feature_name,
                        "operator": predicate_obj.op,
                        "threshold": predicate_obj.threshold,
                        "actual_value": predicate_obj.value,
                    }
                )
            label_info = {
                "probability": row[label + "_proba"],
                "features": pp_models[label][1],
                "predicates": predicates,
            }
            labels_info[label] = label_info
        pp_labels = labels_info

    return pp_labels


def _predict_models(
    m_model, task_type: str, m_meta_features_test: pd.DataFrame, preprocessor_labels: PipelineSkeleton
) -> PipelineSkeleton:
    meta_features_test = m_meta_features_test[[x for x in m_meta_features_test.columns if x.startswith("feature:")]]
    meta_features = meta_features_test.fillna(0)

    predict_proba = m_model[0].predict_proba(meta_features) + m_model[1].predict_proba(meta_features)

    index_ranking = predict_proba.argsort(axis=1)[:, ::-1]
    model_ranking = np.array(m_model[0].classes_)[index_ranking]

    model_ranking = [
        [
            (
                name_to_label_mapping[item][task_type]
                if item in name_to_label_mapping and task_type in name_to_label_mapping[item]
                else None
            )
            for item in preds
        ]
        for preds, task_type in zip(model_ranking, task_type)
    ]
    columns = ["Pred " + str(i + 1) for i in range(len(model_ranking[0]))]
    predictions_df = pd.DataFrame(model_ranking, columns=columns)

    import copy

    all_labels = copy.deepcopy(preprocessor_labels)
    for index, row in predictions_df.iterrows():
        labels = collections.OrderedDict()
        for i, column in enumerate(columns):
            labels[row[column]] = predict_proba[index][index_ranking[index][i]]
        if "ILLEGAL PREDICTION" in labels:
            del labels["ILLEGAL PREDICTION"]

        if None in labels:
            del labels[None]

        all_labels.update(labels)
    return all_labels


def dict_to_df(d) -> pd.DataFrame:
    """Convert dict to DataFrame.

    Parameters
    ----------
    d : dict
        Input dict.

    Returns
    ----------
    pd.DataFrame
        Conversion result.

    """
    return pd.DataFrame.from_dict({0: d}, orient="index", columns=d.keys())


def predict(task: Task, dataset_summary: DatasetSummary) -> PipelineSkeleton:
    """Perform preprocessing and model prediction internal functions.

    Parameters
    ----------
    task : Task
    dataset_summary: DatasetSummary

    Returns
    ----------
    all_labels : PipelineSkeleton
    """
    task_type = task.task_type
    p_meta_feature_test = dict_to_df(dataset_summary.meta_features_pp)
    m_meta_feature_test = dict_to_df(dataset_summary.meta_features_m)
    # check python version and store as a variable
    python_minor_version = sys.version_info.minor

    # Load all the pkl files based on python version
    if python_minor_version in [9, 10, 11]:
        base_path = Path(os.path.dirname(__file__)) / ("../models/PY3" + str(python_minor_version))
        with open(base_path / "pp_models.pkl", "rb") as f:
            pp_model = pickle.load(f)

        with open(base_path / "mp_model_1.pkl", "rb") as f1:
            with open(base_path / "mp_model_2.pkl", "rb") as f2:
                m_model = (pickle.load(f1), pickle.load(f2))

    else:  # Default
        with open(Path(os.path.dirname(__file__)) / "../models/pp_models.pkl", "rb") as f:
            pp_model = pickle.load(f)

        with open(Path(os.path.dirname(__file__)) / "../models/mp_model_1.pkl", "rb") as f1:
            with open(Path(os.path.dirname(__file__)) / "../models/mp_model_2.pkl", "rb") as f2:
                m_model = (pickle.load(f1), pickle.load(f2))
    preprocessor_labels = _predict_preprocessors(pp_model, p_meta_feature_test, search_space.target_labels)
    all_labels = _predict_models(m_model, task_type, m_meta_feature_test, preprocessor_labels)

    return all_labels
