# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.
"""
Train and eval functions used in main.py
"""
import math
import sys
from typing import Iterable, Optional

import torch

import timm
from timm.data import Mixup
from timm.utils import accuracy, ModelEma

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

import numpy as np

from evaluate_model import load_weights, compute_challenge_metric


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, args = None):
    
    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)
        
        outputs = model(samples.float(), if_random_cls_token_position=args.if_random_cls_token_position, if_random_token_rank=args.if_random_token_rank)
        loss = criterion(outputs, targets.float())
            
        loss_value = loss.item()

        if args.lrschedule == "Noam":
            optimizer.optimizer.zero_grad()
        elif args.lrschedule == "CosineAnnealing":
            optimizer.zero_grad()
        else:
            raise ValueError(f"No matching condition for value in the file of engine_ecg_2021.py : {args.lrschedule}")
        
        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)
        if args.lrschedule == "Noam":
            metric_logger.update(lr=optimizer.optimizer.param_groups[0]["lr"])
        elif args.lrschedule == "CosineAnnealing":
            metric_logger.update(lr=optimizer.param_groups[0]["lr"])
        else:
            raise ValueError(f"No matching condition for value in the file of engine_ecg_2021.py for updating : {args.lrschedule}")
        
        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 = []
    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)

        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)
    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)

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

@torch.no_grad()
def evaluate(data_loader, model, thrs_chall, thrs_F1, thrs_accuracy, thrs_hammingLoss, 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)

        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_F1 else 0) for prob in probs] for probs in np.array(outputs)])
    f1 = 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 np.array(outputs)])
    hamming = hamming_loss(targets, outputs_hammingLoss)

    outputs_accuracy = np.array([[(1 if prob > thrs_accuracy else 0) for prob in probs] for probs in np.array(outputs)])
    subset_accuracy = 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 np.array(outputs)])

    challenge_value = compute_challenge_metric(weights, targets, outputs_best, classes, sinus_rhythm_ID)
    print("This is the challenge score:", challenge_value)

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

    return auprc, auroc, f1, hamming, subset_accuracy, challenge_value, loss_value_after_each_epcoh