from typing import Iterable
import torch

from losses import DistillationLoss
import utils

from sklearn.metrics import average_precision_score, roc_auc_score, f1_score
from sklearn.metrics import hamming_loss
from sklearn.metrics import accuracy_score

from sklearn.metrics import recall_score
from sklearn.metrics import precision_score

import numpy as np

from evaluate_model import load_weights, compute_challenge_metric, multilabel_specificity_score
import sys


def normalize_model_outputs(model_outputs):
    """
    Normalize model outputs to the range [0, 1].
    
    Parameters:
    model_outputs (numpy.ndarray): The raw output of the deep learning model.
    
    Returns:
    numpy.ndarray: The normalized model outputs.
    """
    a = model_outputs.min()
    b = model_outputs.max()
    return (model_outputs - a) / (b - a)


sinus_rhythm_ID = set(['426783006'])

weights_file = "./weights.csv"

classes, weights = load_weights(weights_file)

def train_one_epoch(model: torch.nn.Module, criterion: DistillationLoss,
                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
                    device: torch.device, epoch: int,
                    set_training_mode=True):
    
    model.train(set_training_mode)
    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    header = 'Epoch: [{}]'.format(epoch)
    print_freq = 600
    
    output_list = []
    target_list = []    

    loss_value_after_each_epcoh = 0
    batch_num = 0

    for samples, targets in metric_logger.log_every(data_loader, print_freq, header):        
        batch_num += 1

        samples = samples.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)
        samples = samples.unsqueeze(1)

        outputs = model(samples.float())
        loss = criterion(outputs, targets.float())
            
        loss_value = loss.item()
        
        optimizer.zero_grad()
        
        loss.backward() # Backward pass: Compute gradient of the loss with respect to model parameters
            
        # Update parameters/using the Noam
        optimizer.step() 

        torch.cuda.synchronize()
        
        metric_logger.update(loss=loss_value)
        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
        
        target_list.append(targets.data.cpu().numpy())
        output_list.append(outputs.data.cpu().numpy())
        loss_value_after_each_epcoh += loss_value
    
    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)

    # below code is for record the AUPRC of training
    targets_all = np.concatenate(target_list, axis=0)
    outputs_all = np.concatenate(output_list, axis=0)
    outputs_all = normalize_model_outputs(outputs_all)

    threshold = 0.5
    targets_all[targets_all >= threshold] = 1
    targets_all[targets_all < threshold] = 0

    # When using the function below, y_true must be a binarized value.
    train_auprc = average_precision_score(y_true = targets_all, y_score = outputs_all)
    print("This is the training AUPRC:", train_auprc)

    ### The code below is for obtaining the thresholds
    scores_challengeScore = []

    scores_F1 = []
    Macro_scores_F1 = []

    scores_SubsetAccuracy = []
    scores_HammingLoss = []

    for thr in np.arange(0., 1., 0.02):
        outputs_dyn = np.array([[(1 if prob > thr else 0) for prob in probs] for probs in np.array(outputs_all)])

        challenge_value = compute_challenge_metric(weights, targets_all, outputs_dyn, classes, sinus_rhythm_ID)
        scores_challengeScore.append(challenge_value)

        f1 = f1_score(targets_all, outputs_dyn, average='weighted')
        scores_F1.append(f1)

        f1_macro = f1_score(targets_all, outputs_dyn, average='macro')
        Macro_scores_F1.append(f1_macro)

        subset_accuracy = accuracy_score(targets_all, outputs_dyn)
        scores_SubsetAccuracy.append(subset_accuracy)

        hamming = hamming_loss(targets_all, outputs_dyn)
        scores_HammingLoss.append(hamming)

    scores_challengeScore = np.array(scores_challengeScore)
    scores_F1 = np.array(scores_F1)
    ## Macro F1-score
    Macro_scores_F1 = np.array(Macro_scores_F1)

    scores_SubsetAccuracy = np.array(scores_SubsetAccuracy)
    scores_HammingLoss = np.array(scores_HammingLoss)

    # print("This is the challenge score list from training set:\n", scores)

    # Best thrs for challenge score
    thrs_CHALL = np.array([np.argmax(scores_challengeScore, axis=0)*0.02])
    print("This is the best threshold for the challenge score from training set", thrs_CHALL)
    outputs_best_CHALL = np.array([[(1 if prob > thrs_CHALL else 0) for prob in probs] for probs in np.array(outputs_all)])
    challenge_value = compute_challenge_metric(weights, targets_all, outputs_best_CHALL, classes, sinus_rhythm_ID)
    print("This is the challenge score from training set:", challenge_value)

    # Best thrs for F1
    scores_F1 = np.array([np.argmax(scores_F1, axis=0)*0.02])
    print("This is the best threshold for the F1 from training set", scores_F1)
    outputs_best_f1 = np.array([[(1 if prob > scores_F1 else 0) for prob in probs] for probs in np.array(outputs_all)])
    f1 = f1_score(targets_all, outputs_best_f1, average='weighted')
    print("This is the f1 score from training set:", challenge_value)

    # Best thrs for subset accuracy
    scores_SubsetAccuracy = np.array([np.argmax(scores_SubsetAccuracy, axis=0)*0.02])
    print("This is the best threshold for the Subset Accuracy from training set", scores_SubsetAccuracy)
    outputs_best_SubsetAccuracy = np.array([[(1 if prob > scores_SubsetAccuracy else 0) for prob in probs] for probs in np.array(outputs_all)])
    subset_accuracy = accuracy_score(targets_all, outputs_best_SubsetAccuracy)
    print("This is the subset accuracy from training set:", subset_accuracy)

    # Best thrs for hamming loss, here is the loss value from output, we need to find the minimum value.
    # Determine the optimal threshold for Hamming loss. The loss values are obtained from the output, and the goal is to find the minimum value.
    scores_HammingLoss = np.array([np.argmin(scores_HammingLoss, axis=0)*0.02])
    print("This is the best threshold for the Subset Accuracy from training set", scores_HammingLoss)
    outputs_best_HammingLoss = np.array([[(1 if prob > scores_HammingLoss else 0) for prob in probs] for probs in np.array(outputs_all)])
    hamming = hamming_loss(targets_all, outputs_best_HammingLoss)
    print("This is the Hamming from training set:", hamming)

    # Best thrs for Macro F1-score
    Macro_scores_F1 = np.array([np.argmax(Macro_scores_F1, axis=0)*0.02])
    print("This is the best threshold for the Macro F1 from training set", Macro_scores_F1)

    loss_value_after_each_epcoh /= batch_num
    return train_auprc, loss_value_after_each_epcoh, thrs_CHALL, scores_F1, scores_SubsetAccuracy, scores_HammingLoss, Macro_scores_F1

@torch.no_grad()
def evaluate(data_loader, model, thrs_chall, thrs_weighted_F1, thrs_accuracy, thrs_hammingLoss, thrs_Macro_scores_F1, device):
    # criterion = torch.nn.CrossEntropyLoss()
    criterion = torch.nn.BCEWithLogitsLoss()

    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Test:'

    targets = []
    outputs = []
    loss_value_after_each_epcoh = 0
    batch_num = 0
    # switch to evaluation mode
    model.eval()

    for images, target in metric_logger.log_every(data_loader, 100, header):
        
        batch_num += 1
        images = images.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)

        images = images.unsqueeze(1)

        output = model(images.float())

        loss = criterion(output, target.float())
        
        metric_logger.update(loss=loss.item())

        targets.append(target.data.cpu().numpy())
        outputs.append(output.data.cpu().numpy())
    
        loss_value_after_each_epcoh += loss.item()

    # below code is for record the AUPRC of validation
    targets = np.concatenate(targets, axis=0)
    outputs = np.concatenate(outputs, axis=0)
    
    outputs = normalize_model_outputs(outputs)

    auprc = average_precision_score(y_true=targets, y_score=outputs)
    auroc = roc_auc_score(targets, outputs)

    # print("This is the top 3 row:", outputs[0:3])
    outputs_F1 = np.array([[(1 if prob > thrs_weighted_F1 else 0) for prob in probs] for probs in outputs])
    weighted_f1_Tt = f1_score(targets, outputs_F1, average='weighted')

    outputs_hammingLoss = np.array([[(1 if prob > thrs_hammingLoss else 0) for prob in probs] for probs in outputs])
    hamming_Tt = hamming_loss(targets, outputs_hammingLoss)

    outputs_accuracy = np.array([[(1 if prob > thrs_accuracy else 0) for prob in probs] for probs in outputs])
    subset_accuracy_Tt = accuracy_score(targets, outputs_accuracy)

    # print("This is the best threshold for the challenge score, obtained from the training test, without any testing data leakage:", thrs_chall)
    outputs_best = np.array([[(1 if prob > thrs_chall else 0) for prob in probs] for probs in outputs])
    challenge_value_Tt = compute_challenge_metric(weights, targets, outputs_best, classes, sinus_rhythm_ID)
    # print("This is the challenge score:", challenge_value)

    # Marco F-1 score
    outputs_Macro_scores_F1 = np.array([[(1 if prob > thrs_Macro_scores_F1 else 0) for prob in probs] for probs in outputs])
    Macro_scores_F1_Tt = f1_score(targets, outputs_Macro_scores_F1, average='macro')

    # Macro specificity, this threshold is based on the Macro F1-score
    outputs_Macro_specificity = np.array([[(1 if prob > thrs_Macro_scores_F1 else 0) for prob in probs] for probs in outputs])
    Macro_specificity_Tt,_ = multilabel_specificity_score(targets, outputs_Macro_specificity)

    # Weighted specificity, this threshold is based on the Weighted F1-score
    outputs_Weighted_specificity = np.array([[(1 if prob > thrs_weighted_F1 else 0) for prob in probs] for probs in outputs])
    _, Weighted_specificity_Tt= multilabel_specificity_score(targets, outputs_Weighted_specificity)

    # Macro sensitivity, this threshold is based on the Macro F1-score
    outputs_Macro_sensitivity = np.array([[(1 if prob > thrs_Macro_scores_F1 else 0) for prob in probs] for probs in outputs])
    Macro_sensitivity_Tt = recall_score(targets, outputs_Macro_sensitivity, average='macro')

    # Weighted sensitivity, this threshold is based on the Weighted F1-score
    outputs_Weighted_sensitivity = np.array([[(1 if prob > thrs_weighted_F1 else 0) for prob in probs] for probs in outputs])
    Weighted_sensitivity_Tt = recall_score(targets, outputs_Weighted_sensitivity, average='weighted')

    # thrs_Macro_Precision, this threshold is based on the Macro F1-score
    outputs_Macro_Precision = np.array([[(1 if prob > thrs_Macro_scores_F1 else 0) for prob in probs] for probs in outputs])
    Macro_Precision_Tt = precision_score(targets, outputs_Macro_Precision, average='macro', zero_division=1.0)

    # thrs_Weighted_Precision, this threshold is based on the Weighted F1-score
    outputs_Weighted_Precision = np.array([[(1 if prob > thrs_weighted_F1 else 0) for prob in probs] for probs in outputs])
    Weighted_Precision_Tt = precision_score(targets, outputs_Weighted_Precision, average='weighted', zero_division=1.0)

    ### The 5 Best-Value Lists Based on the Six Evaluation Metrics
    scores_challengeScore_list = []

    scores_F1_list = []
    Macro_scores_F1_list = []


    scores_SubsetAccuracy_list = []
    scores_HammingLoss_list = []
    
    #### This is the best value achieved for the threshold-dependent evaluation metrics.
    for thr_b in np.arange(0., 1., 0.02):
        outputs_dyn = np.array([[(1 if prob > thr_b else 0) for prob in probs] for probs in outputs])

        challenge_value = compute_challenge_metric(weights, targets, outputs_dyn, classes, sinus_rhythm_ID)
        scores_challengeScore_list.append(challenge_value)

        ##### for the F1-score(weighted and Macro)
        f1 = f1_score(targets, outputs_dyn, average='weighted')
        scores_F1_list.append(f1)
        f1_macro = f1_score(targets, outputs_dyn, average='macro')
        Macro_scores_F1_list.append(f1_macro)

        subset_accuracy = accuracy_score(targets, outputs_dyn)
        scores_SubsetAccuracy_list.append(subset_accuracy)

        hamming = hamming_loss(targets, outputs_dyn)
        scores_HammingLoss_list.append(hamming)
    
    scores_challengeScore_list = np.array(scores_challengeScore_list)
    best_scores_challengeScore = np.max(scores_challengeScore_list)

    scores_F1_list = np.array(scores_F1_list)
    best_scores_F1 = np.max(scores_F1_list)

    Macro_scores_F1_list = np.array(Macro_scores_F1_list)
    best_Macro_scores_F1 = np.max(Macro_scores_F1_list)

    scores_SubsetAccuracy_list = np.array(scores_SubsetAccuracy_list)
    best_scores_SubsetAccuracy = np.max(scores_SubsetAccuracy_list)
    
    scores_HammingLoss_list = np.array(scores_HammingLoss_list)
    best_scores_HammingLoss = np.min(scores_HammingLoss_list)

    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    loss_value_after_each_epcoh /= batch_num

    return (auprc, auroc, weighted_f1_Tt, hamming_Tt, subset_accuracy_Tt, challenge_value_Tt, Macro_scores_F1_Tt,
            Macro_specificity_Tt, Weighted_specificity_Tt, Macro_sensitivity_Tt, Weighted_sensitivity_Tt, Macro_Precision_Tt,
            Weighted_Precision_Tt, best_scores_challengeScore, best_scores_F1, best_Macro_scores_F1, best_scores_SubsetAccuracy,
            best_scores_HammingLoss, loss_value_after_each_epcoh)