import argparse
import inspect
import os
import sys

# import warnings
from datetime import datetime

import numpy as np
import pandas as pd
from joblib import delayed
from sklearn.base import clone
from sklearn.utils import check_random_state
from sksurv.util import Surv
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 LateFusionSurvival
from multipit.utils.custom.cv import CensoredKFold


def main(params):
    """
    Repeated cross-validation experiment for survival prediction with late fusion
    """

    # Uncomment for disabling ConvergenceWarning
    # warnings.simplefilter("ignore")
    # os.environ["PYTHONWARNINGS"] = 'ignore'

    # 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=True, integration="late")
    ptask.load_data()
    X = ptask.data_concat.values
    y = Surv().from_arrays(
        event=ptask.labels.loc[ptask.data_concat.index, "event"].values,
        time=ptask.labels.loc[ptask.data_concat.index, "time"].values,
    )
    ptask.init_pipelines_latefusion()

    # 3. Perform repeated cross-validation
    parallel = ProgressParallel(
        n_jobs=config["parallelization"]["n_jobs_repeats"],
        total=config["latefusion"]["n_repeats"],
    )
    results_parallel = parallel(
        delayed(_fun_repeats)(
            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"])
    )

    # 4. Save results
    if config["permutation_test"]:
        perm_predictions = np.zeros(
            (
                len(y),
                len(ptask.names),
                config["n_permutations"],
                config["latefusion"]["n_repeats"],
            )
        )
        list_data_preds, list_data_thrs = [], []
        for p, res in enumerate(results_parallel):
            list_data_preds.append(res[0])
            list_data_thrs.append(res[1])
            perm_predictions[:, :, :, p] = res[2]
        perm_labels = results_parallel[-1][3]
        np.save(os.path.join(save_dir, "permutation_labels.npy"), perm_labels)
        np.save(os.path.join(save_dir, "permutation_predictions.npy"), perm_predictions)
        data_preds = pd.concat(list_data_preds, axis=0)
        data_preds.to_csv(os.path.join(save_dir, "predictions.csv"))
        if config["collect_thresholds"]:
            data_thrs = pd.concat(list_data_thrs, axis=0)
            data_thrs.to_csv(os.path.join(save_dir, "thresholds.csv"))
    else:
        list_data_preds, list_data_thrs = [], []
        for p, res in enumerate(results_parallel):
            list_data_preds.append(res[0])
            list_data_thrs.append(res[1])
        data_preds = pd.concat(list_data_preds, axis=0)
        data_preds.to_csv(os.path.join(save_dir, "predictions.csv"))
        if config["collect_thresholds"]:
            data_thrs = pd.concat(list_data_thrs, axis=0)
            data_thrs.to_csv(os.path.join(save_dir, "thresholds.csv"))


def _fun_repeats(prediction_task, X, y, r, disable_infos):
    """
    Train and test a late fusion model for survival task with cross-validation

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

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

    y: Structured array of size (n_samples,)
        Event indicator and observed time for each sample

    r: int
        Repeat number

    disable_infos: bool

    Returns
    -------
    df_pred: pd.DataFrame of shape (n_samples, n_models+4)
        Predictions collected over the test sets of the cross-validation scheme for each multimodal combination

    df_thrs: pd.DataFrame of shape (n_samples, n_models+2), None
        Thresholds that optimize the log-rank test on the training set for each fold and each multimodal combination.

    permut_predictions: 3D array of shape (n_samples, n_models, n_permutations)
        Predictions collected over the test sets of the cross_validation scheme for each multimodal combination and each random permutation of the labels.

    permut_labels: 3D array of shape (n_samples, n_permutations, 2)
        Permuted event indicators and observed times.
    """
    cv = CensoredKFold(n_splits=10, shuffle=True)  # , random_state=np.random.seed(i))
    X_preds = np.zeros((len(y), 4 + len(prediction_task.names)))
    X_thresholds = (
        np.zeros((len(y), 2 + len(prediction_task.names)))
        if prediction_task.config["collect_thresholds"]
        else None
    )
    late_clf = LateFusionSurvival(
        estimators=prediction_task.late_estimators,
        cv=CensoredKFold(n_splits=10, shuffle=True, random_state=np.random.seed(r)),
        **prediction_task.config["latefusion"]["args"]
    )

    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],
        )

        clf = clone(late_clf)
        clf.fit(X_train, y_train)

        for c, idx in enumerate(prediction_task.indices):
            X_preds[test_index, c] = clf.predict(X_test, estim_ind=idx)
            if prediction_task.config["collect_thresholds"]:
                X_thresholds[test_index, c] = clf.find_logrank_threshold(
                    X_train, y_train, estim_ind=idx
                )

        X_preds[test_index, -4] = fold_index
        if prediction_task.config["collect_thresholds"]:
            X_thresholds[test_index, -2] = fold_index

    X_preds[:, -3] = r
    if prediction_task.config["collect_thresholds"]:
        X_thresholds[:, -1] = r
    X_preds[:, -2] = y["time"]
    X_preds[:, -1] = y["event"]
    df_pred = (
        pd.DataFrame(
            X_preds,
            columns=prediction_task.names
            + ["fold_index", "repeat", "label.time", "label.event"],
            index=prediction_task.data_concat.index,
        )
        .reset_index()
        .rename(columns={"index": "samples"})
        .set_index(["repeat", "samples"])
    )

    if prediction_task.config["collect_thresholds"]:
        df_thrs = (
            pd.DataFrame(
                X_thresholds,
                columns=prediction_task.names + ["fold_index", "repeat"],
                index=prediction_task.data_concat.index,
            )
            .reset_index()
            .rename(columns={"index": "samples"})
            .set_index(["repeat", "samples"])
        )
    else:
        df_thrs = None

    permut_predictions = None
    permut_labels = None
    if prediction_task.config["permutation_test"]:
        permut_labels = np.zeros((len(y), prediction_task.config["n_permutations"], 2))
        permut_predictions = np.zeros(
            (
                len(y),
                len(prediction_task.names),
                prediction_task.config["n_permutations"],
            )
        )
        for prm in range(prediction_task.config["n_permutations"]):
            X_perm = np.zeros((len(y), len(prediction_task.names)))
            random_state = check_random_state(prm)
            sh_ind = random_state.permutation(len(y))
            yshuffle = np.copy(y)[sh_ind]
            permut_labels[:, prm, 0] = yshuffle["time"]
            permut_labels[:, prm, 1] = yshuffle["event"]
            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, :],
                    yshuffle[train_index],
                    X[test_index, :],
                    yshuffle[test_index],
                )
                clf = clone(late_clf)
                clf.fit(X_train, y_train)

                for c, idx in enumerate(prediction_task.indices):
                    X_perm[test_index, c] = clf.predict(X_test, estim_ind=idx)
            permut_predictions[:, :, prm] = X_perm
    return df_pred, df_thrs, permut_predictions, permut_labels


if __name__ == "__main__":
    args = argparse.ArgumentParser(description="Late fusion")
    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())