from sklearn.model_selection import train_test_split

import pandas as pd

import numpy as np

import torch

from joblib import load

import statistics as stats

from sklearn import preprocessing

import torch.backends.cudnn as cudnn

cudnn.enabled = True

cudnn.benchmark = False

cudnn.deterministic = True

from code_psd_shallow_eeg_gcnn.EEGGraphDataset import EEGGraphDataset

from code_psd_shallow_eeg_gcnn.EEGGraphConvNet import EEGGraphConvNet

from torch_geometric.data import DataLoader

from torch.utils.data import WeightedRandomSampler

from sklearn.metrics import make_scorer

from sklearn.metrics import balanced_accuracy_score, auc, accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, roc_curve

from torchvision.transforms import Compose, ToTensor

stats_test_data = { }

# after each epoch, record all the metrics on both train and validation sets

def collect_metrics(y_probs_test, y_true_test, y_pred_test, sample_indices_test,

                    fold_idx, experiment_name):

    dataset_index = pd.read_csv("master_metadata_index.csv", dtype={"patient_ID":str, })

    # create patient-level train and test dataframes

    rows = [ ]

    for i in range(len(sample_indices_test)):

        idx = sample_indices_test[i]

        temp = { }

        temp["patient_ID"] = str(dataset_index.loc[idx, "patient_ID"])

        temp["sample_idx"] = idx

        temp["y_true"] = y_true_test[i]

        temp["y_probs_0"] = y_probs_test[i, 0]

        temp["y_probs_1"] = y_probs_test[i, 1]

        temp["y_pred"] = y_pred_test[i]

        rows.append(temp)

    test_patient_df = pd.DataFrame(rows)

    # get patient-level metrics from window-level dataframes

    y_probs_test_patient, y_true_test_patient, y_pred_test_patient = get_patient_prediction(test_patient_df, fold_idx)

    stats_test_data[f"probs_0_fold_{fold_idx}"] = y_probs_test_patient[:, 0]

    stats_test_data[f"probs_1_fold_{fold_idx}"] = y_probs_test_patient[:, 1]

    window_csv_dict = { }

    patient_csv_dict = { }

    # WINDOW-LEVEL ROC PLOT

    # pos_label="healthy"

    fpr, tpr, thresholds = roc_curve(y_true_test, y_probs_test[:,1], pos_label=1)

    window_csv_dict[f"fpr_fold_{fold_idx}"] = fpr

    window_csv_dict[f"tpr_fold_{fold_idx}"] = tpr

    window_csv_dict[f"thres_fold_{fold_idx}"] = thresholds

    # PATIENT-LEVEL ROC PLOT - select optimal threshold for this, and get patient-level precision, recall, f1

    # pos_label="healthy"

    fpr, tpr, thresholds = roc_curve(y_true_test_patient, y_probs_test_patient[:,1], pos_label=1)

    patient_csv_dict[f"fpr_fold_{fold_idx}"] = fpr

    patient_csv_dict[f"tpr_fold_{fold_idx}"] = tpr

    patient_csv_dict[f"thres_fold_{fold_idx}"] = thresholds

    # select an optimal threshold using the ROC curve

    # Youden's J statistic to obtain the optimal probability threshold and this method gives equal weights to both false positives and false negatives

    optimal_proba_cutoff = sorted(list(zip(np.abs(tpr - fpr), thresholds)), key=lambda i: i[0], reverse=True)[0][1]

    # print (optimal_proba_cutoff)

    # calculate class predictions and confusion-based metrics using the optimal threshold

    roc_predictions = [1 if i >= optimal_proba_cutoff else 0 for i in y_probs_test_patient[:,1]]

    precision_patient_test =  precision_score(y_true_test_patient, roc_predictions, pos_label=0)

    recall_patient_test =  recall_score(y_true_test_patient, roc_predictions, pos_label=0)

    f1_patient_test = f1_score(y_true_test_patient, roc_predictions, pos_label=0)

    bal_acc_patient_test = balanced_accuracy_score(y_true_test_patient, roc_predictions)

    # PATIENT-LEVEL AUROC

    from sklearn.metrics import roc_auc_score

    auroc_patient_test = roc_auc_score(y_true_test_patient, y_probs_test_patient[:,1])

    # AUROC

    from sklearn.metrics import roc_auc_score

    # CAUTION - The binary case expects a shape (n_samples,), and the scores must be the scores of the class with the greater label.

    # https://scikit-learn.org/stable/modules/generated/sklearn.metrics.roc_auc_score.html

    auroc_test = roc_auc_score(y_true_test, y_probs_test[:,1])

    return auroc_patient_test, auroc_test, precision_patient_test, recall_patient_test, f1_patient_test, bal_acc_patient_test

# create patient-level metrics

def get_patient_prediction(df, fold_idx):

    unique_patients = list(df["patient_ID"].unique())

    grouped_df = df.groupby("patient_ID")

    rows = [ ]

    for patient in unique_patients:

        patient_df = grouped_df.get_group(patient)

        temp = { }

        temp["patient_ID"] = patient

        temp["y_true"] = list(patient_df["y_true"].unique())[0]

        assert len(list(patient_df["y_true"].unique())) == 1

        temp["y_pred"] = patient_df["y_pred"].mode()[0]

        temp["y_probs_0"] = patient_df["y_probs_0"].mean()

        temp["y_probs_1"] = patient_df["y_probs_1"].mean()

        rows.append(temp)

    return_df = pd.DataFrame(rows)

    # need subject names and labels for comparisons testing

    if fold_idx == 0:

        stats_test_data["subject_id"] = list(return_df["patient_ID"][:])

        stats_test_data["label"] = return_df["y_true"][:]

    return np.array(list(zip(return_df["y_probs_0"], return_df["y_probs_1"]))), list(return_df["y_true"]), list(return_df["y_pred"])

if __name__ == "__main__":

    GPU_IDX = 0

    EXPERIMENT_NAME = "psd_gnn_shallow"

    BATCH_SIZE = 512

    SFREQ = 250.0

    NUM_EPOCHS = 100

    NUM_WORKERS = 6

    PIN_MEMORY = True

    # ensure reproducibility of results

    SEED = 42

    np.random.seed(SEED)

    torch.manual_seed(SEED)

    print("[MAIN] Numpy and PyTorch seed set to {} for reproducibility.".format(SEED))

    MASTER_DATASET_INDEX = pd.read_csv("master_metadata_index.csv", dtype={"patient_ID":str, })

    subjects = MASTER_DATASET_INDEX["patient_ID"].astype("str").unique()

    print("[MAIN] Subject list fetched! Total subjects are {}...".format(len(subjects)))

    # NOTE: splitting whole subjects into train+validation and heldout test

    train_val_subjects, test_subjects = train_test_split(subjects, test_size=0.30, random_state=SEED)

    print("[MAIN] (Train + validation) and (heldout test) split made at subject level. 30 percent subjects held out for testing.")  

    train_subjects, val_subjects = train_test_split(train_val_subjects, test_size=0.20, random_state=SEED)

    train_indices = MASTER_DATASET_INDEX.index[MASTER_DATASET_INDEX["patient_ID"].astype("str").isin(train_subjects)].tolist()

    val_indices = MASTER_DATASET_INDEX.index[MASTER_DATASET_INDEX["patient_ID"].astype("str").isin(val_subjects)].tolist()

    # use GPU when available

    DEVICE = torch.device('cuda:{}'.format(GPU_IDX) if torch.cuda.is_available() else 'cpu')

    torch.cuda.set_device(DEVICE)

    print('[MAIN] Using device:', DEVICE, torch.cuda.get_device_name(DEVICE))

    X = load("psd_features_data_X")

    y = load("labels_y")

    # normalize psd_features_data_X

    normd_x = []

    for i in range(len(y)):

        arr = X[i, :]

        arr = arr.reshape(1, -1)

        arr2 = preprocessing.normalize(arr)

        arr2 = arr2.reshape(48)

        normd_x.append(arr2)

    norm = np.array(normd_x)

    X = norm.reshape(len(y), 48)

    # get 0/1 labels for pytorch, ensure mapping is the same between train and test

    label_mapping, y = np.unique(y, return_inverse = True)

    print("[MAIN] unique labels to [0 1] mapping:", label_mapping)

    model = EEGGraphConvNet(reduced_sensors=False)

    model = model.to(DEVICE).double()

    labels_unique, counts = np.unique(y, return_counts=True)

    class_weights = np.array([1.0/x for x in counts])

    # provide weights for samples in the training set only      

    sample_weights = class_weights[y[train_indices]]

    # sampler needs to come up with training set size number of samples

    weighted_sampler = WeightedRandomSampler(weights=sample_weights, num_samples=len(train_indices), replacement=True)

    # define training set

    train_dataset = EEGGraphDataset(X=X, y=y, indices=train_indices, loader_type="train", 

                                    sfreq=SFREQ, transform=Compose([ToTensor()]))

    train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, sampler=weighted_sampler,

                             num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY)

    # define validation set

    val_dataset = EEGGraphDataset(X=X, y=y, indices=val_indices, loader_type="validation", 

                                    sfreq=SFREQ, transform=Compose([ToTensor()]))

    val_loader = DataLoader(dataset=val_dataset, batch_size=BATCH_SIZE, 

                              shuffle=False, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY)

    # define loss function

    loss_function = torch.nn.CrossEntropyLoss()

    # define optimizer

    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    # define scheduler

    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[i*10 for i in range(1, 26)], gamma=0.1)

    # start training

    for epoch in range(NUM_EPOCHS):

        model.train()

        train_loss = []

        val_loss = []

        y_probs_train = torch.empty(0, 2).to(DEVICE)

        y_true_train = [ ]

        y_pred_train = [ ]

        window_indices_train = [ ]

        for batch_idx, batch in enumerate(train_loader):

            # send batch to GPU

            X_batch = batch.to(device=DEVICE, non_blocking=True)

            y_batch = torch.tensor(batch.y)

            y_batch = y_batch.to(device=DEVICE, non_blocking=True)

            window_indices_train += X_batch.dataset_idx.cpu().numpy().tolist()

            optimizer.zero_grad()

            # forward pass

            outputs = model(X_batch.x, X_batch.edge_index, X_batch.edge_attr, X_batch.batch).float()

            loss = loss_function(outputs, y_batch)

            train_loss.append(loss.item())

            # backward pass

            loss.backward()

            _, predicted = torch.max(outputs.data, 1)

            y_pred_train += predicted.cpu().numpy().tolist()

            # concatenate along 0th dimension

            y_probs_train = torch.cat((y_probs_train, outputs.data), 0)

            y_true_train += y_batch.cpu().numpy().tolist()

            optimizer.step()

        scheduler.step()

        # returning prob distribution over target classes, take softmax across the 1st dimension

        y_probs_train = torch.nn.functional.softmax(y_probs_train, dim=1).cpu().numpy()

        y_true_train = np.array(y_true_train)

        # calculate training set metrics

        auroc_patient_train, auroc_train, precision_patient_train, recall_patient_train, f1_patient_train, bal_acc_patient_train = collect_metrics(y_probs_test=y_probs_train,

                        y_true_test=y_true_train,

                        y_pred_test=y_pred_train,

                        sample_indices_test = window_indices_train,                 

                        fold_idx=0,

                        experiment_name=EXPERIMENT_NAME)

        # evaluate on validation set

        model.eval()

        with torch.no_grad():

            y_probs_val = torch.empty(0, 2).to(DEVICE)

            y_true_val = [ ]

            y_pred_val = [ ]

            window_indices_val = [ ]

            for i, batch in enumerate(val_loader):

                X_batch = batch.to(device=DEVICE, non_blocking=True)

                y_batch = torch.tensor(batch.y)

                y_batch = y_batch.to(device=DEVICE, non_blocking=True)

                window_indices_val += X_batch.dataset_idx.cpu().numpy().tolist()

                outputs = model(X_batch.x, X_batch.edge_index, X_batch.edge_attr, X_batch.batch).float()

                loss = loss_function(outputs, y_batch)

                val_loss.append(loss.item())

                _, predicted = torch.max(outputs.data, 1)

                y_pred_val += predicted.cpu().numpy().tolist()

                # concatenate along 0th dimension

                y_probs_val = torch.cat((y_probs_val, outputs.data), 0)

                y_true_val += y_batch.cpu().numpy().tolist()

        # returning prob distribution over target classes, take softmax across the 1st dimension

        y_probs_val = torch.nn.functional.softmax(y_probs_val, dim=1).cpu().numpy()

        y_true_val = np.array(y_true_val)

        # get validation set metrics

        auroc_patient_val, auroc_val, precision_patient_val, recall_patient_val, f1_patient_val, bal_acc_patient_val = collect_metrics(y_probs_test=y_probs_val,

                        y_true_test=y_true_val,

                        y_pred_test=y_pred_val,

                        sample_indices_test = val_indices,                  

                        fold_idx=0,

                        experiment_name=EXPERIMENT_NAME)

        # save the model every 20 epochs

        if epoch % 20 == 0:

            state = {

                'model_description': str(model),

                'state_dict': model.state_dict(),

                'optimizer': optimizer.state_dict()

            }

            torch.save(state, f"model_{epoch}.ckpt")

        print(f'Epoch: {epoch}-----------------------------------------------------------')

        print(f"Train loss: {np.mean(train_loss):.3f}; Validation loss: {np.mean(val_loss):.3f}")

        print(f"Train AUROC:{auroc_train:.3f}; Validation AUROC: {auroc_val:.3f}")

        print(f"Train patient metrics: AUROC{auroc_patient_train:.3f}, precision: {precision_patient_train:.3f}, recall: {recall_patient_train:.3f}, f1: {f1_patient_train:.3f}, bal acc: {bal_acc_patient_train:.3f}")

        print(f"Validation patient metrics: AUROC{auroc_patient_val:.3f}, precision: {precision_patient_val:.3f}, recall: {recall_patient_val:.3f}, f1: {f1_patient_val:.3f}, bal acc: {bal_acc_patient_val:.3f}")