import torch

import os

from utils import get_scheduler, get_model, get_criterion, get_optimizer, get_metric

from data_functions import get_loaders

from utils import OneHotEncoder

import wandb

from tqdm import tqdm

import time

# helping function to normal visualisation in Colaboratory

def foo_():

    time.sleep(0.3)

def train_epoch(model, train_dl, encoder, criterion, metric, optimizer, scheduler, device):

    model.train()

    loss_sum = 0

    score_sum = 0

    with tqdm(total=len(train_dl), position=0, leave=True) as pbar:

        for X, y in tqdm(train_dl, position=0, leave=True):

            pbar.update()

            X = X.to(device)

            if len(torch.unique(X)) == 1:

                continue

            if encoder is not None:

                y = encoder(y)

            y = y.squeeze(4)

            y = y.to(device)

            optimizer.zero_grad()

            output = model(X)

            loss = criterion(output, y)

            loss.backward()

            optimizer.step()

            scheduler.step()

            loss = loss.item()

            score = metric(output, y).mean().item()

            loss_sum += loss

            score_sum += score

    return loss_sum / len(train_dl), score_sum / len(train_dl)

def eval_epoch(model, val_dl, encoder, criterion, metric, device):

    model.eval()

    loss_sum = 0

    score_sum = 0

    with tqdm(total=len(val_dl), position=0, leave=True) as pbar:

        for X, y in tqdm(val_dl, position=0, leave=True):

            pbar.update()

            X = X.to(device)

            if len(torch.unique(X)) == 1:

                continue

            if encoder is not None:

                y = encoder(y)

            y = y.squeeze()

            y = y.to(device)

            with torch.no_grad():

                output = model(X)

                loss = criterion(output, y).item()

                score = metric(output, y).mean().item()

                loss_sum += loss

                score_sum += score

    return loss_sum / len(val_dl), score_sum / len(val_dl)

def run(cfg, model_name, use_wandb=True, max_early_stopping=2):

    torch.cuda.empty_cache()

    # <<<<< SETUP >>>>>

    train_loader, val_loader = get_loaders(cfg)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    model = get_model(cfg)(cfg=cfg).to(device)

    optimizer = get_optimizer(cfg)(model.parameters(), **cfg.optimizer_params)

    scheduler = get_scheduler(cfg)(optimizer, **cfg.scheduler_params)

    metric = get_metric(cfg)(**cfg.metric_params)

    criterion = get_criterion(cfg)(**cfg.criterion_params)

    encoder = OneHotEncoder(cfg)

    # wandb is watching

    if use_wandb:

        wandb.init(project='Covid19_CT_segmentation_' + str(cfg.dataset_name), entity='aiijcteamname', config=cfg,

                   name=model_name)

        wandb.watch(model, log_freq=100)

    best_val_loss = 999

    last_train_loss = 0

    last_val_loss = 999

    early_stopping_flag = 0

    best_state_dict = model.state_dict()

    for epoch in range(1, cfg.epochs + 1):

        print(f'Epoch #{epoch}')

        # <<<<< TRAIN >>>>>

        train_loss, train_score = train_epoch(model, train_loader, encoder,

                                              criterion, metric,

                                              optimizer, scheduler, device)

        print('      Score    |    Loss')

        print(f'Train: {train_score:.6f} | {train_loss:.6f}')

        # <<<<< EVAL >>>>>

        val_loss, val_score = eval_epoch(model, val_loader, encoder,

                                         criterion, metric, device)

        print(f'Val: {val_score:.6f} | {val_loss:.6f}', end='\n\n')

        metrics = {'train_score': train_score,

                   'train_loss': train_loss,

                   'val_score': val_score,

                   'val_loss': val_loss,

                   'lr': scheduler.get_last_lr()[-1]}

        if use_wandb:  # log metrics to wandb

            wandb.log(metrics)

        # saving best weights

        if val_loss < best_val_loss:

            best_val_loss = val_loss

            best_state_dict = model.state_dict()

            torch.save(best_state_dict, os.path.join('checkpoints', model_name + '.pth'))

        # weapon counter over-fitting

        if train_loss < last_train_loss and val_loss > last_val_loss:

            early_stopping_flag += 1

        if early_stopping_flag == max_early_stopping:

            print('<<< EarlyStopping >>>')

            break

        last_train_loss = train_loss

        last_val_loss = val_loss

    # loading best weights

    model.load_state_dict(best_state_dict)

    if use_wandb:

        wandb.finish()

    return model