import argparse
import inspect
import os
import sys

# import warnings
from datetime import datetime

import numpy as np
import pandas as pd
import shap
from joblib import delayed
from sklearn.base import clone
from sklearn.model_selection import StratifiedKFold
from tqdm import tqdm

from _init_scripts import PredictionTask
from _utils import read_yaml, write_yaml, ProgressParallel

currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(currentdir)
sys.path.insert(0, parentdir)

from multipit.multi_model.latefusion import LateFusionClassifier


def main(params):
    """ """

    # 0. Read config file and save it in the results
    config = read_yaml(params.config)
    save_name = config["save_name"]
    if save_name is None:
        run_id = datetime.now().strftime(r"%m%d_%H%M%S")
        save_name = "exp_" + run_id
    save_dir = os.path.join(params.save_path, save_name)
    os.mkdir(save_dir)
    write_yaml(config, os.path.join(save_dir, "config.yaml"))

    # 1. fix random seeds for reproducibility
    seed = config["latefusion"]["seed"]
    np.random.seed(seed)

    # 2. Load data and define pipelines for each modality
    ptask = PredictionTask(config, survival=False, integration="late")
    ptask.load_data()
    X, y = ptask.data_concat.values, ptask.labels.loc[ptask.data_concat.index].values
    ptask.init_pipelines_latefusion()

    parallel = ProgressParallel(
        n_jobs=config["parallelization"]["n_jobs_repeats"],
        total=config["latefusion"]["n_repeats"],
    )
    list_shap = parallel(
        delayed(_fun_parallel)(
            ptask,
            X,
            y,
            r,
            disable_infos=(config["parallelization"]["n_jobs_repeats"] is not None)
            and (config["parallelization"]["n_jobs_repeats"] > 1),
        )
        for r in range(config["latefusion"]["n_repeats"])
    )

    shap_explain = {"clinical": [], "radiomics": [], "pathomics": [], "RNA": []}
    coefs_LR = {"clinical": [], "radiomics": [], "pathomics": [], "RNA": []}

    for results in list_shap:
        for moda, shapley in results[0].items():
            shap_explain[moda].append(shapley)

    for key, val in shap_explain.items():
        df_shap = pd.concat(val, axis=0, join="outer")
        df_shap.to_csv(os.path.join(save_dir, "Shap_" + key + ".csv"))

    if config["classifier"]["type"] == "LR":
        for results in list_shap:
            for moda, coefs in results[1].items():
                coefs_LR[moda].append(coefs)

        for key, val in coefs_LR.items():
            coefficients = np.stack(val, axis=-1)
            np.save(os.path.join(save_dir, "coef_LR_" + key + ".npy"), coefficients)


def _fun_parallel(prediction_task, X, y, r, disable_infos):
    """
    Collect SHAP values for several unimodal classifiers with cross-validation

    Parameters
    ----------
    prediction_task: PredictionTask object

    X: 2D array of shape (n_samples, n_features)
        Concatenation of the different modalities

    y: 1D array of shape (n_samples,)
        Binary outcome

    r: int
        Repeat number

    disable_infos: bool

    Returns
    -------
    shap_dict: dictionary
        Dictionary whose keys correspond to the different modalities (e.g., "RNA", "clinical") and the items correspond
        to pandas dataframe of size (n_samples, n_features) that contain the SHAP values collected across the test sets
        of the cross-validation scheme.

    coefs_dict: dictionary or None
        Dictionary whose keys correspond to the different modalities (e.g., "RNA", "clinical") and the items correspond
        to arrays of size (n_folds, n_features) that contain the linear coefficients collected across the different
        folds of the cross-validation scheme. None if the classifier type is not linear.
    """

    cv = StratifiedKFold(n_splits=10, shuffle=True)
    late_clf = LateFusionClassifier(
        estimators=prediction_task.late_estimators,
        cv=StratifiedKFold(n_splits=10, shuffle=True, random_state=np.random.seed(r)),
        **prediction_task.config["latefusion"]["args"]
    )

    shap_dict = {name: [] for name, *_ in late_clf.estimators}

    if prediction_task.config["classifier"]["type"] == "LR":
        coef_dict = {name: [] for name, *_ in late_clf.estimators}

    for fold_index, (train_index, test_index) in tqdm(
        enumerate(cv.split(np.zeros(len(y)), y)),
        leave=False,
        total=cv.get_n_splits(np.zeros(len(y))),
        disable=disable_infos,
    ):
        X_train, y_train, X_test, y_test = (
            X[train_index, :],
            y[train_index],
            X[test_index, :],
            y[test_index],
        )
        # Fit late fusion on the training set of the fold
        clf = clone(late_clf)
        clf.fit(X_train, y_train)
        # Collect SHAP values on the test set of the fold for each unimodal classifier
        for ind, (name, estim, features) in enumerate(clf.fitted_estimators_):
            X_background = X_train[:, features]
            bool_mask = ~(
                np.sum(np.isnan(X_background), axis=1)
                > clf.missing_threshold * len(features)
            )
            X_explain = X_test[:, features]
            bool_mask_explain = ~(
                np.sum(np.isnan(X_explain), axis=1)
                > clf.missing_threshold * len(features)
            )
            if clf.calibration is not None:
                explainer = shap.Explainer(
                    lambda x: (
                        clf.fitted_meta_estimators_[(ind,)].predict_proba(
                            estim.predict_proba(x)[:, 1].reshape(-1, 1)
                        )
                    ),
                    X_background[bool_mask, :],
                )
            else:
                explainer = shap.Explainer(
                    lambda x: estim.predict_proba(x), X_background[bool_mask, :]
                )
            shap_values = explainer(X_explain[bool_mask_explain, :])
            shap_df = pd.DataFrame(
                shap_values.values[:, :, 1],
                columns=prediction_task.data_concat.columns[features],
                index=prediction_task.data_concat.index.values[
                    test_index[bool_mask_explain]
                ],
            )
            shap_df["fold_index"] = fold_index
            shap_df["repeat"] = r
            shap_dict[name].append(shap_df)
            # Also collect coefficients for logistic regreression
            if prediction_task.config["classifier"]["type"] == "LR":
                if name == "RNA":
                    temp = np.zeros((1, 40))
                    temp[:, : estim[-1].coef_.shape[1]] = estim[-1].coef_
                    coef_dict[name].append(temp)
                else:
                    coef_dict[name].append(estim[-1].coef_)

    if prediction_task.config["classifier"]["type"] == "LR":
        coefs_dict = {name: np.vstack(value) for name, value in coef_dict.items()}
    else:
        coefs_dict = None

    shap_dict = {
        name: pd.concat(value, axis=0, join="outer")
        for name, value in shap_dict.items()
    }

    return shap_dict, coefs_dict


if __name__ == "__main__":
    args = argparse.ArgumentParser(description="Collect Shap")
    args.add_argument(
        "-c",
        "--config",
        type=str,
        help="config file path",
    )
    args.add_argument(
        "-s",
        "--save_path",
        type=str,
        help="save path",
    )
    main(params=args.parse_args())