import gc
import itertools
import random

import numpy as np
import optuna
import pandas as pd
import torch

from evaluation import multiclass_acc
from model import FakeNewsModel, calculate_loss
from utils import AvgMeter, print_lr, EarlyStopping, CheckpointSaving


def batch_constructor(config, batch):
    b = {}
    for key, value in batch.items():
        if key != 'text':
            b[key] = value.to(config.device)
        else:
            b[key] = value
    return b


def train_epoch(config, model, train_loader, optimizer, scalar):
    loss_meter = AvgMeter('train')
    c_loss_meter = AvgMeter('train')
    s_loss_meter = AvgMeter('train')

    # tqdm_object = tqdm(train_loader, total=len(train_loader))

    targets = []
    predictions = []
    for index, batch in enumerate(train_loader):
        batch = batch_constructor(config, batch)
        optimizer.zero_grad(set_to_none=True)

        with torch.cuda.amp.autocast():
            output, score = model(batch)
            loss, c_loss, s_loss = calculate_loss(model, score, batch['label'])
        scalar.scale(loss).backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)

        if (index + 1) % 2:
            scalar.step(optimizer)
            # loss.backward()
            # optimizer.step()
            scalar.update()

        count = batch["id"].size(0)
        loss_meter.update(loss.detach(), count)
        c_loss_meter.update(c_loss.detach(), count)
        s_loss_meter.update(s_loss.detach(), count)

        # tqdm_object.set_postfix(train_loss=loss_meter.avg, lr=get_lr(optimizer))

        prediction = output.detach()
        predictions.append(prediction)

        target = batch['label'].detach()
        targets.append(target)

    losses = (loss_meter, s_loss_meter, c_loss_meter)

    return losses, targets, predictions


def validation_epoch(config, model, validation_loader):
    loss_meter = AvgMeter('validation')
    c_loss_meter = AvgMeter('validation')
    s_loss_meter = AvgMeter('validation')

    targets = []
    predictions = []
    # tqdm_object = tqdm(validation_loader, total=len(validation_loader))
    for batch in validation_loader:
        batch = batch_constructor(config, batch)
        with torch.no_grad():
            output, score = model(batch)
            loss, c_loss, s_loss = calculate_loss(model, score, batch['label'])

            count = batch["id"].size(0)
            loss_meter.update(loss.detach(), count)
            c_loss_meter.update(c_loss.detach(), count)
            s_loss_meter.update(s_loss.detach(), count)

            # tqdm_object.set_postfix(validation_loss=loss_meter.avg)

            prediction = output.detach()
            predictions.append(prediction)

            target = batch['label'].detach()
            targets.append(target)

    losses = (loss_meter, s_loss_meter, c_loss_meter)

    return losses, targets, predictions


def supervised_train(config, train_loader, validation_loader, trial=None):
    torch.cuda.empty_cache()
    checkpoint_path2 = checkpoint_path = str(config.output_path) + '/checkpoint.pt'
    if trial:
        checkpoint_path2 = str(config.output_path) + '/checkpoint_' + str(trial.number) + '.pt'

    torch.manual_seed(27)
    random.seed(27)
    np.random.seed(27)
    
    torch.autograd.set_detect_anomaly(False)
    torch.autograd.profiler.profile(False)
    torch.autograd.profiler.emit_nvtx(False)

    scalar = torch.cuda.amp.GradScaler()
    model = FakeNewsModel(config).to(config.device)

    params = [
        {"params": model.image_encoder.parameters(), "lr": config.image_encoder_lr, "name": 'image_encoder'},
        {"params": model.text_encoder.parameters(), "lr": config.text_encoder_lr, "name": 'text_encoder'},
        {"params": itertools.chain(model.image_projection.parameters(), model.text_projection.parameters()),
         "lr": config.head_lr, "weight_decay": config.head_weight_decay, 'name': 'projection'},
        {"params": model.classifier.parameters(), "lr": config.classification_lr,
         "weight_decay": config.classification_weight_decay,
         'name': 'classifier'}
    ]
    optimizer = torch.optim.AdamW(params, amsgrad=True)
    lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=config.factor,
                                                              patience=config.patience // 5, verbose=True)
    early_stopping = EarlyStopping(patience=config.patience, delta=config.delta, path=checkpoint_path, verbose=True)
    checkpoint_saving = CheckpointSaving(path=checkpoint_path, verbose=True)

    train_losses, train_accuracies = [], []
    validation_losses, validation_accuracies = [], []

    validation_accuracy, validation_loss = 0, 1
    for epoch in range(config.epochs):
        print(f"Epoch: {epoch + 1}")
        gc.collect()

        model.train()
        train_loss, train_truth, train_pred = train_epoch(config, model, train_loader, optimizer, scalar)
        model.eval()
        with torch.no_grad():
            validation_loss, validation_truth, validation_pred = validation_epoch(config, model, validation_loader)

        train_accuracy = multiclass_acc(train_truth, train_pred)
        validation_accuracy = multiclass_acc(validation_truth, validation_pred)
        print_lr(optimizer)
        print('Training Loss:', train_loss[0], 'Training Accuracy:', train_accuracy)
        print('Validation Loss', validation_loss, 'Validation Accuracy:', validation_accuracy)

        if lr_scheduler:
            lr_scheduler.step(validation_loss[0].avg)
        if early_stopping:
            early_stopping(validation_loss[0].avg, model)
            if early_stopping.early_stop:
                print("Early stopping")
                break
        if checkpoint_saving:
            checkpoint_saving(validation_accuracy, model)

        train_accuracies.append(train_accuracy)
        train_losses.append(train_loss)
        validation_accuracies.append(validation_accuracy)
        validation_losses.append(validation_loss)

        if trial:
            trial.report(validation_accuracy, epoch)
            if trial.should_prune():
                print('trial pruned')
                raise optuna.exceptions.TrialPruned()

        print()

    if checkpoint_saving:
        model = FakeNewsModel(config).to(config.device)
        model.load_state_dict(torch.load(checkpoint_path))
        model.eval()
        with torch.no_grad():
            validation_loss, validation_truth, validation_pred = validation_epoch(config, model, validation_loader)
        validation_accuracy = multiclass_acc(validation_pred, validation_truth)
        if trial and validation_accuracy >= config.wanted_accuracy:
            loss_accuracy = pd.DataFrame(
                {'train_loss': train_losses, 'train_accuracy': train_accuracies, 'validation_loss': validation_losses,
                 'validation_accuracy': validation_accuracies})
            torch.save({'model_state_dict': model.state_dict(),
                        'parameters': str(config),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'loss_accuracy': loss_accuracy}, checkpoint_path2)

    if not checkpoint_saving:
        loss_accuracy = pd.DataFrame(
            {'train_loss': train_losses, 'train_accuracy': train_accuracies, 'validation_loss': validation_losses,
             'validation_accuracy': validation_accuracies})
        torch.save(model.state_dict(), checkpoint_path)
        if trial and validation_accuracy >= config.wanted_accuracy:
            torch.save({'model_state_dict': model.state_dict(),
                        'parameters': str(config),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'loss_accuracy': loss_accuracy}, checkpoint_path2)
    try:
        del train_loss
        del train_truth
        del train_pred
        del validation_loss
        del validation_truth
        del validation_pred
        del train_losses
        del train_accuracies
        del validation_losses
        del validation_accuracies
        del loss_accuracy
        del scalar
        del model
        del params
    except:
        print('Error in deleting caches')
        pass

    return validation_accuracy