"""
Baselines on OpenBHB with classical ML + DL algorithms on:
- 5-fold CV stratified on age+sex+site scheme + Train/Val/Test Split (3 runs)
- VBM, quasi-raw, FreeSurfer pre-processing
- 2 tasks: age regression and sex classification

"""

import numpy as np
import os, logging
import sys, argparse
from copy import deepcopy
sys.path.extend(["../", ".", "../../"])
from pynet.datasets.open_bhb import OpenBHB
from pynet.open_bhb_ml.sk_trainer import OpenBHBMLTrainer
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.svm import SVC, SVR
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.preprocessing import StandardScaler


"""
Execuction of sk-learn models on Train/Val/Test Split + 5-fold CV Split for OpenBHB-10K for:
* VBM + Quasi-Raw (+ FSL?)
* Age + Sex 
* Linear + SVM + Random Forest 

python3 age_sex_prediction.py --preproc $p --scheme $s --pb $pb --model $m --njobs 2 --nb_folds 3 
"""

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--preproc", required=True, choices=["quasi_raw", "vbm"])
    parser.add_argument("--scheme", required=True, choices=["train_val_test", "cv"])
    parser.add_argument("--pb", required=True, choices=["age", "sex"])
    parser.add_argument("--model", required=True, choices=["Ridge", "SVR", "SVC", "LogisticRegression",
                                                           "RandomForestClassifier",
                                                           "RandomForestRegressor"])
    parser.add_argument("--njobs", default=1, type=int)
    parser.add_argument("--folds", nargs='+', type=int)
    parser.add_argument("--nb_folds", type=int)

    args = parser.parse_args()
    root = '/neurospin/psy_sbox/analyses/201906_schizconnect-vip-prague-bsnip-biodb-icaar-start_assemble-all/data/'
    saving_dir = '/neurospin/psy_sbox/bd261576/checkpoints/regression_age_sex/Benchmark_OpenBHB'

    scorings = {"age": "neg_mean_absolute_error",
                "sex": "balanced_accuracy"}

    masks = {"vbm": np.load(os.path.join(root, "mask_open-bhb_vbm.npy")),
             "quasi_raw": np.load(os.path.join(root, "mask_open-bhb_quasi_raw.npy"))}

    models = {"age": [Ridge, SVR, RandomForestRegressor],
              "sex": [SVC, LogisticRegression, RandomForestClassifier]}

    models_hyperparams = {"age": {"Ridge": {'alpha': [1e-1, 1, 10], 'fit_intercept': [True]},
                                  "SVR": {'kernel': ['rbf'], 'C': [1e2, 1e1, 1]},
                                  "RandomForestRegressor": {'n_estimators': [10, 100, 150],
                                                            'min_samples_split': [2, 4],
                                                            'min_samples_leaf': [1, 4],
                                                            'n_jobs': [6]},
                                  },
                          "sex": {"SVC": {'kernel': ['rbf'], 'C': [1e2, 1e1, 1]},
                                  "RandomForestClassifier": {'n_estimators': [10, 100, 150],
                                                             'min_samples_split': [2, 4],
                                                             'min_samples_leaf': [1, 4],
                                                             'n_jobs': [6]},
                                  "LogisticRegression": {'C': [1e2, 1e1, 1], 'solver': ['sag'], 'penalty': ['l2'],
                                                         'n_jobs': [6],
                                                         'fit_intercept': [True]}
                                  }
                          }
    dbs = {"train_val_test": "OpenBHB-Train-Val-Test", "cv": "OpenBHB-CV"}

    logger = logging.getLogger("pynet")

    preproc = args.preproc
    label = args.pb
    pb = args.pb
    model = eval(args.model)
    assert model in models[pb]
    scheme = args.scheme
    db = dbs[scheme]
    scaler = StandardScaler
    hyperparams = models_hyperparams[pb][args.model]
    mask = masks[args.preproc]
    scoring = scorings[pb]

    folds = list(args.folds or range(args.nb_folds))
    folds_per_job = np.array_split(folds, np.ceil(len(folds)/args.njobs))
    # Parallelize across folds
    for folds_job_i in folds_per_job:
        jobs = []
        logger.info("Processing folds {f} - Model {m} - Scheme {s} - Label {l} - Preproc {p}".format(f=folds_job_i,
                                                                                                     m=args.model,
                                                                                                     s=scheme,
                                                                                                     l=label,
                                                                                                     p=preproc))
        for fold in folds_job_i:
            train = OpenBHB(root, preproc=preproc, scheme=scheme, split="train", fold=fold, target=label)
            test = OpenBHB(root, preproc=preproc, scheme=scheme, split="test", fold=fold, target=label)
            model_name = type(model()).__name__
            preproc_name = "" if preproc == "vbm" else "quasi_raw"
            exp_name = 'Test_{m}_{pb}_{db}_fold{k}_epoch{e}.pkl'.format(m=model_name, pb=pb.capitalize(),
                                                                        db=db,
                                                                        k=fold, e=100)
            saving_dir_ = os.path.join(saving_dir, preproc_name, model_name, pb.capitalize())
            trainer = OpenBHBMLTrainer(model(), deepcopy(hyperparams), train, test,
                                       mask=mask, exp_name=exp_name, saving_dir=saving_dir_,
                                       scoring=scoring, scaler=scaler(), n_jobs=5, logger=logger)
            trainer.start()
            jobs.append(trainer)
        for job in jobs: job.join()
