""" Training augmented model """
import os
import torch
import torch.nn as nn
import numpy as np
from tensorboardX import SummaryWriter
from ptflops import get_model_complexity_info
import utils
import data_generator_3D as data_generator_3D
import time
import SimpleITK as sitk
import sys
from config import TrainConfig
from model import LCOVNet
from apex import amp


config = TrainConfig()

device = torch.device("cuda")

# tensorboard
writer = SummaryWriter(log_dir=os.path.join(config.path, "tb"))
writer.add_text('config', config.as_markdown(), 0)

logger = utils.get_logger(os.path.join(config.path, "{}.log".format(config.name)))
config.print_params(logger.info)

def main():
    logger.info("Logger is set - training start")

    # set default gpu device id
    torch.cuda.set_device(config.gpus[0])

    # set seed
    np.random.seed(config.seed)
    torch.manual_seed(config.seed)
    torch.cuda.manual_seed_all(config.seed)

    torch.backends.cudnn.benchmark = True


    criterion = utils.log_loss().to(device)
    d = torch.device(type='cuda', index=config.gpus[0])
    model = LCOVNet(config.input_channels, config.n_classes).to(device=d)
    with torch.cuda.device(config.gpus[0]):
        net = model
        macs, params = get_model_complexity_info(net, (1, 240, 160, 48), as_strings=True,
                                                 print_per_layer_stat=True, verbose=True)
        logger.info("{:<30}  {:<8}".format('Computational complexity: ', macs))
        logger.info("{:<30}  {:<8}".format('Number of parameters: ', params))

    # model size
    mb_params = utils.param_size(model)
    logger.info("Model size = {:.3f} MB".format(mb_params))
    # weights optimizer
    optimizer = torch.optim.SGD(model.parameters(), config.lr, momentum=config.momentum,
                                weight_decay=config.weight_decay)
    model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

    train_loader = data_generator_3D.Covid19TrainSet()
    valid_loader = data_generator_3D.Covid19EvalSet()
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, config.epochs)

    best_dice = 0.
    # training loop
    summ_writer = SummaryWriter(config.training_summary_dir)
    for epoch in range(config.epochs):

        # training
        train(train_loader, model, optimizer, criterion, epoch, summ_writer)
        lr_scheduler.step()
        # validation
        cur_step = (epoch+1) * len(train_loader)
        mean_dice = validate(valid_loader, model, criterion, epoch, summ_writer, best_dice)

        # save
        if best_dice < mean_dice:
            best_dice = mean_dice
            is_best = True
        else:
            is_best = False
        utils.save_checkpoint(model, config.path, is_best)
        print("")

    logger.info("Final best Dice = {:.4%}".format(best_dice))
    utils.save_results(best_dice, config.path)
    summ_writer.close()

def train(train_loader, model, optimizer, criterion, epoch, summ_writer):
    losses = utils.AverageMeter()
    cur_step = epoch*len(train_loader)
    cur_lr = optimizer.param_groups[0]['lr']
    logger.info("Epoch {} LR {}".format(epoch, cur_lr))
    writer.add_scalar('train/lr', cur_lr, cur_step)
    model.train()
    #all_dice = np.empty().astype(np.float32)
    all_dice = []
    for step, (name, X, y) in enumerate(train_loader):
        X, y = torch.from_numpy(X).to(device, non_blocking=True), torch.from_numpy(y).to(device, non_blocking=True)
        N = X.size(0)

        optimizer.zero_grad()
        logits = model(X)

        loss = criterion(logits, y)
        #loss.backward()
        with amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward()
        # gradient clipping
        nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
        optimizer.step()

        losses.update(loss.item(), N)

        if step % config.print_freq == 0 or step == len(train_loader)-1:
           logger.info(
                "Train: [{:2d}/{}] Step {:03d}/{:03d} Loss {:.3f} ".format(
                    epoch+1, config.epochs, step, len(train_loader), losses.avg,
                    ))

        writer.add_scalar('train/loss', loss.item(), cur_step)

        logits[logits >= 0.5] = 1
        logits[logits < 0.5] = 0
        predict = logits.cpu().detach().numpy()
        y = y.cpu().detach().numpy()
        dice_i = utils.evaluate(predict, y)
        all_dice.append(dice_i)
        cur_step += 1
    dice_mean = 0

    for i in all_dice:
        dice_mean += i/len(all_dice)

    train_avg_loss = losses.avg
    train_avg_dice = dice_mean
    loss_scalers = {'train': train_avg_loss}
    summ_writer.add_scalars('loss', loss_scalers, epoch + 1)

    dice_scalers = {'train': train_avg_dice}
    summ_writer.add_scalars('avg_dice', dice_scalers, epoch + 1)

    if (epoch+1) % 50 == 0:
        chpt_prefx = config.training_checkpoint_prefix
        save_dict = {'epoch': epoch + 1,
                     'model_state_dict': model.state_dict(),
                     'optimizer_state_dict': optimizer.state_dict(),
                     'amp': amp.state_dict()}
        save_name = "{0:}_{1:}.pt".format(chpt_prefx, epoch + 1)
        torch.save(save_dict, save_name)
    print("train_avg_loss", train_avg_loss)
    print("train_avg_dice", train_avg_dice)

def validate(valid_loader, model, criterion, epoch, summ_writer, best_dice):
    losses = utils.AverageMeter()

    model.eval()
    all_dice = np.zeros([len(valid_loader)]).astype(np.float32)
    all_dice = []
    totel_time = 0
    start_time = time.time()
    size_z = 48
    with torch.no_grad():
        for i, (name, image, label) in enumerate(valid_loader):
            image = torch.from_numpy(image)
            predict = np.zeros(shape=label.shape, dtype=label.dtype)
            z = image.shape[4]
            m = z // size_z if z % size_z == 0 else z // size_z + 1
            start_time = time.time()
            for k in range(m):
                if (k+1)*size_z <= z:
                    max_z = (k+1)*size_z
                else:
                    max_z = z
                min_z = max_z - size_z
                image_k = image[:, :, :, :, min_z:max_z].float().to(device, non_blocking=True)
                predict_k = model(image_k)
                predict_k[predict_k >= 0.5] = 1
                predict_k[predict_k < 0.5] = 0
                predict[:, :, :, :, min_z:max_z] = predict_k.cpu().detach().numpy()
            totel_time = totel_time + time.time() - start_time
            all_dice.append(utils.evaluate(predict, label))

    dice_len = len(all_dice)
    dice_np = np.empty(shape=[dice_len])
    #list_image = []
    for i in range(dice_len):
        dice_np[i] = all_dice[i]
        logger.info("{}  dice: {:.4%} ".format(i, all_dice[i]))
    logger.info("mean: {}".format(dice_np.mean()))
    logger.info("std : {}".format(dice_np.std()))

    if best_dice < dice_np.mean():
        chpt_prefx = config.validing_checkpoint_prefix
        save_dict = {'epoch': epoch + 1,
                     'model_state_dict': model.state_dict(),
                     'amp': amp.state_dict()}
        fname = "{}/best.pt".format(chpt_prefx)
        if os.path.isfile(fname):
            os.remove(fname)
        save_name = "{}/best.pt".format(chpt_prefx)
        torch.save(save_dict, save_name)

    dice_scalers = {'vadil': dice_np.mean()}
    summ_writer.add_scalars('vadil_avg_dice', dice_scalers, epoch + 1)

    avg_time = totel_time / dice_len
    logger.info("average testing time : {}".format(avg_time))

    mean_dice = np.mean(all_dice, axis = 0)
    writer.add_scalar('val/dice', mean_dice, epoch)
    writer.add_scalar('val/loss', losses.avg, epoch)
    logger.info("Valid: [{:2d}/{}] average dice: {:.4%} ".format(epoch+1, config.epochs, mean_dice))

    return mean_dice



def save_nd_array_as_image(data, image_name, reference_name = None):
    """
    save a 3D or 2D numpy array as medical image or RGB image
    inputs:
        data: a numpy array with shape [D, H, W] or [C, H, W]
        image_name: the output file name
    outputs: None
    """
    data_dim = len(data.shape)
    assert(data_dim == 2 or data_dim == 3)
    if (image_name.endswith(".nii.gz") or image_name.endswith(".nii") or
        image_name.endswith(".mha")):
        assert(data_dim == 3)
        save_array_as_nifty_volume(data, image_name, reference_name)

def save_array_as_nifty_volume(data, image_name, reference_name = None):
    """
    save a numpy array as nifty image
    inputs:
        data: a numpy array with shape [Depth, Height, Width]
        image_name: the ouput file name
        reference_name: file name of the reference image of which affine and header are used
    outputs: None
    """
    img = sitk.GetImageFromArray(data)
    if(reference_name is not None):
        img_ref = sitk.ReadImage(reference_name)
        img.CopyInformation(img_ref)
    sitk.WriteImage(img, image_name)


if __name__ == "__main__":
    main()