import argparse

import inspect

import os

import sys

from datetime import datetime

import numpy as np

import pandas as pd

from joblib import delayed

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.earlyfusion import EarlyFusionClassifier

def main(params):

    """

    Repeated cross-validation experiment for classification with late fusion

    """

    # 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["earlyfusion"]["seed"]

    np.random.seed(seed)

    config = read_yaml(params.config)

    # 2. Load data and define pipeline

    ptask = PredictionTask(config, survival=False, integration="early")

    ptask.load_data()

    X, y = ptask.data_concat.values, ptask.labels.loc[ptask.data_concat.index].values

    ptask.init_pipelines_earlyfusion()

    # 3. Perform repeated cross-validation

    parallel = ProgressParallel(

        n_jobs=config["parallelization"]["n_jobs_repeats"],

        total=config["earlyfusion"]["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["earlyfusion"]["n_repeats"])

    )

    # 4. Save results

    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 an early fusion model for classification 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

    -------

    df_pred: pd.DataFrame of shape (n_samples, n_models+3)

        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.

    """

    cv = StratifiedKFold(n_splits=10, shuffle=True)

    X_preds = np.zeros((len(y), 3 + len(prediction_task.names)))

    X_thresholds = (

        np.zeros((len(y), 2 + len(prediction_task.names)))

        if prediction_task.config["collect_thresholds"]

        else None

    )

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

        )

        target_surv_train = prediction_task.target_surv[train_index]

        cv_inner = StratifiedKFold(

            n_splits=10, shuffle=True, random_state=np.random.seed(r)

        )

        for c, models in enumerate(prediction_task.names):

            t = {

                model: prediction_task.early_transformers[model]

                for model in models.split("+")

            }

            early_clf = EarlyFusionClassifier(

                estimator=prediction_task.early_estimator,

                transformers=t,

                modalities={

                    model: prediction_task.dic_modalities[model]

                    for model in models.split("+")

                },

                cv=cv_inner,

                **prediction_task.config["earlyfusion"]["classifier_args"]

            )

            if len(models.split("+")) == 1:

                early_clf.set_params(**{"select_features": False})

            early_clf.fit(X_train, y_train)

            X_preds[test_index, c] = early_clf.predict_proba(X_test)[:, 1]

            if prediction_task.config["collect_thresholds"]:

                X_thresholds[test_index, c] = early_clf.find_logrank_threshold(

                    X_train, target_surv_train

                )

        X_preds[test_index, -3] = fold_index

        if prediction_task.config["collect_thresholds"]:

            X_thresholds[test_index, -2] = fold_index

    X_preds[:, -2] = r

    if prediction_task.config["collect_thresholds"]:

        X_thresholds[:, -1] = r

    X_preds[:, -1] = y

    df_pred = (

        pd.DataFrame(

            X_preds,

            columns=prediction_task.names + ["fold_index", "repeat", "label"],

            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

    return df_pred, df_thrs

if __name__ == "__main__":

    args = argparse.ArgumentParser(description="Early 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())