import torch
import matplotlib.pyplot as plt
import numpy as np

from tqdm import tqdm
import os
import random
from torch.cuda.amp import autocast

def set_multiprosessing():
    import torch.multiprocessing
    torch.multiprocessing.set_sharing_strategy('file_system')

def set_seeds(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False  # for faster training, but not deterministic


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def calculate_topk_accuracy(y_pred, y, k=2):
    with torch.no_grad():
        batch_size = y.shape[0]
        _, top_pred = y_pred.topk(k, 1)
        top_pred = top_pred.t()
        correct = top_pred.eq(y.view(1, -1).expand_as(top_pred))
        correct_1 = correct[:1].reshape(-1).float().sum(0, keepdim=True)
        correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
        acc_1 = correct_1 / batch_size
        acc_k = correct_k / batch_size
    return acc_1, acc_k


def train(model, iterator, optimizer, criterion, scaler, device):
    epoch_loss = 0
    epoch_acc_1 = 0
    epoch_acc_5 = 0

    model.train()
    tot = 0
    correct = 0
    for (x, y) in tqdm(iterator):
        x, y = x.cuda(), y.cuda()
        with autocast(enabled=True):
            optimizer.zero_grad()
            y_pred = model(x)

            loss = criterion(y_pred, y)
            tot += y.size(0)
            acc_1, acc_5 = calculate_topk_accuracy(y_pred, y)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            _, predicted = y_pred.max(1)
            correct += predicted.eq(y).sum().item()
            epoch_loss += loss.item()
            epoch_acc_1 += acc_1.item()
            epoch_acc_5 += acc_5.item()
    print("Accuracy = {}".format(100. * correct / tot))

    epoch_loss /= len(iterator)
    epoch_acc_1 /= len(iterator)
    epoch_acc_5 /= len(iterator)

    return epoch_loss, epoch_acc_1, epoch_acc_5


def evaluate(model, iterator, criterion, device):
    epoch_loss = 0
    epoch_acc_1 = 0
    epoch_acc_5 = 0

    model.eval()
    with torch.no_grad():
        for (x, y) in tqdm(iterator):
            x, y = x.cuda(), y.cuda()
            y_pred = model(x)
            loss = criterion(y_pred, y)

            acc_1, acc_5 = calculate_topk_accuracy(y_pred, y)

            epoch_loss += loss.item()
            epoch_acc_1 += acc_1.item()
            epoch_acc_5 += acc_5.item()

    epoch_loss /= len(iterator)
    epoch_acc_1 /= len(iterator)
    epoch_acc_5 /= len(iterator)

    return epoch_loss, epoch_acc_1, epoch_acc_5


def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs



def train_ensemble(model, iterator, optimizer, criterion, scaler, device):
    epoch_loss = 0
    epoch_acc_1 = 0
    epoch_acc_5 = 0

    model.train()
    tot = 0
    correct = 0
    for (x, y) in tqdm(iterator):
        x, y = x.cuda(), y.cuda()
        with autocast(enabled=True):
            optimizer.zero_grad()
            y_pred = model(x)

            loss = criterion(y_pred, y)
            tot += y.size(0)
            acc_1, acc_5 = calculate_topk_accuracy(y_pred, y)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            _, predicted = y_pred.max(1)
            correct += predicted.eq(y).sum().item()
            epoch_loss += loss.item()
            epoch_acc_1 += acc_1.item()
            epoch_acc_5 += acc_5.item()
    print("Accuracy = {}".format(100. * correct / tot))

    epoch_loss /= len(iterator)
    epoch_acc_1 /= len(iterator)
    epoch_acc_5 /= len(iterator)

    return epoch_loss, epoch_acc_1, epoch_acc_5


def evaluate_ensemble(model, iterator, criterion, device):
    epoch_loss = 0
    epoch_acc_1 = 0
    epoch_acc_5 = 0

    model.eval()
    with torch.no_grad():
        for (x, y) in tqdm(iterator):
            x, y = x.cuda(), y.cuda()
            y_pred = model(x)
            loss = criterion(y_pred, y)

            acc_1, acc_5 = calculate_topk_accuracy(y_pred, y)

            epoch_loss += loss.item()
            epoch_acc_1 += acc_1.item()
            epoch_acc_5 += acc_5.item()

    epoch_loss /= len(iterator)
    epoch_acc_1 /= len(iterator)
    epoch_acc_5 /= len(iterator)

    return epoch_loss, epoch_acc_1, epoch_acc_5


def AUROC_curve(true_labels, pred_probs):
    import sklearn.metrics
    fpr, tpr, thresholds = sklearn.metrics.roc_curve(y_true=true_labels, y_score=pred_probs,
                                                     pos_label=1)  # positive class is 1; negative class is 0
    auroc = sklearn.metrics.auc(fpr, tpr)


