"""
@author: gbello & lisuru6
How to run the code
python demo_validateDL.py -c /path-to-conf

Default conf uses demo/scripts/default_validate_DL.conf

"""
import json
import shutil
from datetime import timedelta
import pickle
import numpy as np
from pathlib import Path
from argparse import ArgumentParser
from lifelines.utils import concordance_index

from survival4D.nn import hypersearch_nn
from survival4D.nn import train_nn
from survival4D.config import NNExperimentConfig, HypersearchConfig, ModelConfig
from matplotlib import pyplot as plt
from sklearn.model_selection import KFold

DEFAULT_CONF_PATH = Path(__file__).parent.joinpath("default_nn.conf")


def parse_args():
    parser = ArgumentParser()
    parser.add_argument(
        "-c", "--conf-path", dest="conf_path", type=str, default=None, help="Conf path."
    )
    return parser.parse_args()


def main():
    args = parse_args()
    if args.conf_path is None:
        conf_path = DEFAULT_CONF_PATH
    else:
        conf_path = Path(args.conf_path)
    exp_config = NNExperimentConfig.from_conf(conf_path)
    exp_config.output_dir.mkdir(parents=True, exist_ok=True)
    hypersearch_config = HypersearchConfig.from_conf(conf_path)
    model_config = ModelConfig.from_conf(conf_path)

    shutil.copy(str(conf_path), str(exp_config.output_dir.joinpath("nn.conf")))

    # import input data: i_full=list of patient IDs, y_full=censoring status and survival times for patients,
    # x_full=input data for patients (i.e. motion descriptors [11,514-element vector])

    with open(str(exp_config.data_path), 'rb') as f:
        c3 = pickle.load(f)
    x_full = c3[0]
    y_full = c3[1]
    print(x_full.shape, y_full.shape)

    del c3

    # Initialize lists to store predictions
    c_vals = []
    c_trains = []

    kf = KFold(n_splits=exp_config.n_folds)
    i = 0
    for train_indices, test_indices in kf.split(x_full):

        x_train, y_train = x_full[train_indices], y_full[train_indices]
        x_val, y_val = x_full[test_indices], y_full[test_indices]

        # STEP 1: find optimal hyperparameters using CV
        print("Step 1a")
        opars, osummary = hypersearch_nn(
            x_data=x_train,
            y_data=y_train,
            method=exp_config.search_method,
            nfolds=exp_config.n_folds,
            nevals=exp_config.n_evals,
            batch_size=exp_config.batch_size,
            num_epochs=exp_config.n_epochs,
            backend=exp_config.backend,
            model_kwargs=model_config.to_dict(),
            **hypersearch_config.to_dict(),
        )
        print("Step b")
        # (1b) using optimal hyperparameters, train a model and test its performance on the holdout validation set.
        olog = train_nn(
            backend=exp_config.backend,
            xtr=x_train,
            ytr=y_train,
            batch_size=exp_config.batch_size,
            n_epochs=exp_config.n_epochs,
            **model_config.to_dict(),
            **opars,
        )

        # (1c) Compute Harrell's Concordance index
        pred_val = olog.predict(x_val, batch_size=1)[1]
        c_val = concordance_index(y_val[:, 1], -pred_val, y_val[:, 0])

        pred_train = olog.predict(x_train, batch_size=1)[1]
        c_train = concordance_index(y_train[:, 1], -pred_train, y_train[:, 0])
        c_vals.append(c_val)
        c_trains.append(c_train)
        save_params(
            opars, osummary, "cv_{}".format(i), exp_config.output_dir,
            c_val=c_val, c_train=c_train,
            c_val_mean=np.mean(c_vals), c_val_var=np.var(c_vals),
            c_train_mean=np.mean(c_trains), c_train_var=np.var(c_trains)
        )
        print('Validation concordance index = {0:.4f}'.format(c_val))
        i += 1
        plot_cs(c_trains, c_vals, exp_config.output_dir)
    print('Mean Validation concordance index = {0:.4f}'.format(np.mean(c_vals)))
    print('Variance = {0:.4f}'.format(np.var(c_vals)))


def save_params(params: dict, search_log, name: str, output_dir: Path, **kwargs):
    output_dir.mkdir(parents=True, exist_ok=True)
    params["search_log_optimum_c_index"] = search_log.optimum
    params["num_evals"] = search_log.stats["num_evals"]
    params["time"] = str(timedelta(seconds=search_log.stats["time"]))
    params["call_log"] = search_log.call_log
    for key in kwargs.keys():
        params[key] = kwargs[key]
    with open(str(output_dir.joinpath(name + ".json")), "w") as fp:
        json.dump(params, fp, indent=4)


def compute_bootstrap_adjusted_c_index(C_app, Cb_opts):
    # Compute bootstrap-estimated optimism (mean of optimism estimates across the B bootstrap samples)
    C_opt = np.mean(Cb_opts)

    # Adjust apparent C using bootstrap-estimated optimism
    C_adj = C_app - C_opt

    # compute confidence intervals for optimism-adjusted C
    C_opt_95confint = np.percentile([C_app - o for o in Cb_opts], q=[2.5, 97.5])

    return C_opt, C_adj, C_opt_95confint


def plot_cs(c_trains, c_vals, output_dir):
    plt.figure()
    plt.title("CV validation, mean={:.4f}, var={:.4f}".format(np.mean(c_vals), np.var(c_vals)))
    plt.plot(range(len(c_vals)), c_vals, 'rx-')

    plt.plot(range(len(c_trains)), c_trains, 'bx-')
    plt.savefig(str(output_dir.joinpath("c_train_val.png")))


if __name__ == '__main__':
    main()