import argparse
import os
import numpy as np
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
from brats import UncertainBraTS
from probabilistic_unet import ProbabilisticUnet
import hydra
from omegaconf import DictConfig, OmegaConf
from utils import TensorboardSummary
from saver import Saver
from metrics import Evaluator

from utils import l2_regularisation
from torch.utils.data.sampler import SubsetRandomSampler


class qubiqTrainer(object):
    # Define Saver
    def __init__(self, args):
        self.args = args

         # Define Dataloader
        kwargs = {'num_workers': args.training.workers, 'pin_memory': True}
       


        train_dataset = UncertainBraTS(mode = 'train', dataset= args.tasks.dataset, task = args.tasks.task, output= 'individual')
        val_dataset = UncertainBraTS(mode = 'val', dataset = args.tasks.dataset, task = args.tasks.task, output= 'individual' )


        self.train_length = len(train_dataset)
        train_indices = list(range(self.train_length))
        np.random.shuffle(train_indices)
        sampler = SubsetRandomSampler(train_indices)

        self.train_loader = DataLoader(train_dataset, batch_size = args.training.train_size, sampler = sampler, **kwargs)
        self.val_loader = DataLoader(val_dataset, batch_size = args.training.val_size, shuffle = False, **kwargs)
        self.nclass = 1
        self.train_length = len(train_dataset)
        print('number of classes: ', self.nclass)

        # Define the parameters for the sample evaluation
        self.num_sample = args.training.num_sample

        # Define network
        model = ProbabilisticUnet(input_channels=args.tasks.inchannel, num_classes=self.nclass, 
        num_filters=args.training.filters, latent_dim=args.training.latent_dim, 
        no_convs_fcomb=args.training.no_convs_fcomb, 
        beta=args.training.beta)
        
        # Define Save
        self.saver = Saver(args)
        self.saver.save_experiment_config()

        # Define Tensorboard Summary
        self.summary = TensorboardSummary(self.saver.experiment_dir)
        self.writer = self.summary.create_summary()
    

        # Define number of epochs
        self.num_epoch = args.training.epoch


        # set up the learning rate
        train_params = [{'params': model.parameters(), 'lr': args.training.lr}]

        # Define Optimizer
        optimizer = torch.optim.SGD(train_params, momentum=args.training.momentum,
                                weight_decay=args.training.weight_decay, 
                                nesterov=args.training.nesterov)

        self.optimizer = optimizer
        # Define lr scheduler
        # self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
        #                           args.epochs, len(self.train_loader))

        # Model 
        self.model = model        
        # Define Evaluator
        self.evaluator = Evaluator(self.nclass, dice=True, loss='ce')

        # Using cuda
        if args.training.cuda:
            # self.model = torch.nn.DataParallel(self.model, device_ids=self.args.gpu_ids)
            self.model = self.model.cuda()


    def training(self, epoch):
        train_loss = 0.0
        kl_loss = 0.0

        self.model.train()
        tbar = tqdm(self.train_loader)
        num_img_tr = len(self.train_loader)
        for i, sample in enumerate(tbar):
            image, target = sample['image'], sample['label']
            if self.args.training.cuda:
                image, target = image.cuda(), target.cuda()

            self.optimizer.zero_grad()
            # print(image.shape)
            # print(target.shape)
            # target = torch.unsqueeze(target, 1)
            self.model.forward(image, target, training=True)
            elbo = self.model.elbo(target)
            reg_loss = l2_regularisation(self.model.posterior) + l2_regularisation(self.model.prior) + l2_regularisation(self.model.fcomb.layers)
            loss = -elbo + self.args.training.alpha * reg_loss
            loss.backward()
            self.optimizer.step()
            train_loss += loss.item()
            kl_loss += reg_loss.item()

            tbar.set_description('Train loss: %.4f' % (train_loss / (i + 1)))
            # tbar.set_description("Train kl loss: %.4f" % (kl_loss / (i + 1)))

            self.writer.add_scalar('train/total_loss_iter', loss.item(), i + num_img_tr * epoch)
            self.writer.add_scalar("train/total_reg_loss_iter", kl_loss, i + num_img_tr * epoch)

        # Show 10 * 3 inference results each epoch
        global_step = i + num_img_tr * epoch
        self.summary.visualize(self.writer, image, target, output, global_step)

        self.writer.add_scalar('train/total_loss_epoch', train_loss / i, epoch)
        self.writer.add_scalar("train/total_kl_loss_epoch", kl_loss / i, epoch)
        print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size + image.data.shape[0]))
        print('Loss: %.3f' % (train_loss))
        print("KL: %.4f" % (kl_loss))

        if self.args.saver.no_val:
            # save checkpoint every epoch
            is_best = False
            self.saver.save_checkpoint({
                'epoch': epoch + 1,
                'state_dict': self.model.module.state_dict(),
                'optimizer': self.optimizer.state_dict(),
                'best_pred': self.best_pred,
            }, is_best)

    def val(self, epoch):
        self.model.eval()
        self.evaluator.reset()
        tbar = tqdm(self.val_loader, desc='\r')
        test_loss = 0.0
        kl_loss = 0.0
        ncc_score = 0.0
        for i, sample in enumerate(tbar):
            image, target = sample['image'], sample['label']

            if self.args.cuda:
                image, target = image.cuda(), target.cuda()
            with torch.no_grad():
                output = self.model.forward(image, target, training=False)

            # tbar.set_description("Test KL Loss: %.4f" % (kl_loss / (i + 1)))

            pred = output.data.cpu().numpy()
            target = target.data.cpu().numpy()

            self.evaluator.add_batch(target, pred)

        qubiq_score = self.evaluator.QUBIQ_score()
        ged = self.evaluator.GED()
        sd = self.evaluator.SD()
        sa = self.evaluator.SA()

        self.writer.add_scalar('QUBIQ score', qubiq_score, epoch)
        self.writer.add_scalar("NCC score", ncc_score, epoch)
        self.writer.add_scalar("GED score", ged, epoch)
        self.writer.add_scalar("Sample diversity", sd, epoch)
        self.writer.add_scalar("Sample accuracy", sa, epoch)

        print('Validation:')
        print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size + image.data.shape[0]))
        # print("dice: {}".format(dice))
        # print("Shape of dice_class: {}".format(dice_class.shape))
        print("QUBIQ score {}".format(qubiq_score))
        print("NCC score {}".format(ncc_score))
        print("GED score {}".format(ged))
        print("Sample diversity {}".format(sd))
        print("Sample accuracy {}".format(sa))
        print('Loss: %.3f' % (test_loss))

        is_best = True
        self.best_pred = qubiq_score
        self.saver.save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': self.model.module.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            'best_pred': self.best_pred,
        }, is_best)

    def test(self, path):
        """
        For QUBIQ dataset, there are no test labels.
        This founction is only for lidc dataset.
        :return:
        """
        checkpoint = torch.load(path, map_location=torch.device("cuda" if self.args.cuda else "cpu"))
        if self.args.cuda:
            self.model.module.load_state_dict(checkpoint["state_dict"], strict=True)
        else:
            self.model.load_state_dict(checkpoint['state_dict'], strict = True)

        self.model.eval()
        self.evaluator.reset()
        tbar = tqdm(self.test_loader, desc='\r')
        test_loss = 0.0
        ged_list = []

        for i, sample in enumerate(tbar):
            image, target = sample['image'], sample['label']
            n, c, w, h = target.shape
            predictions = target.data.new(self.num_sample, n, c, w, h)
            if self.args.cuda:
                image, target = image.cuda(), target.cuda()
            kl_losses = torch.zeros((self.num_sample, n))
            with torch.no_grad():
                for j in range(self.num_sample):
                    output, kl = self.model(image)
                    # print(output.shape)
                    if self.args.cuda:

                        predictions[j] = output.cpu()
                        kl_losses[j] = kl.cpu()
                    else:
                        predictions[j] = output
                        kl_losses[j] = kl

            mean_out = torch.mean(predictions.float(), dim=0, keepdim=False)
            mean_kl_loss = torch.mean(kl_losses.float())

            if self.args.cuda:
                test_loss = metrics.dice_coef(mean_out, target.cpu(), self.nclass)
            else:
                test_loss = metrics.dice_coef(mean_out, target, self.nclass)

            pred = mean_out.data.cpu().numpy()
            target = target.data.cpu().numpy()

            self.evaluator.add_batch(target, pred)

            tbar.set_description('Sample Dice loss: %.3f' % (test_loss / (i + 1)))
            tbar.set_description("Sample KL Loss: %.4f" % (mean_kl_loss / (i + 1)))

        qubiq_score = self.evaluator.QUBIQ_score()
        sd = self.evaluator.SD()
        sa = self.evaluator.SA()
        ged = np.mean(ged_list)

        print('Sampling %d times for Test:' % (self.num_sample))
        print('[numImages: %5d]' % (i * self.args.batch_size + image.data.shape[0]))

        print("Sampling QUBIQ score {}".format(qubiq_score))
        print("Sampling GED score {}".format(ged))
        print("Sampling diversity {}".format(sd))
        print("Sampling accuracy {}".format(sa))
        print('Sampling Dice: %.3f' % (test_loss))

@hydra.main(config_path="conf", config_name="config")
def main(cfg: DictConfig):

    # default settings for epochs, batch_size and lr

    print(OmegaConf.to_yaml(cfg))
    torch.cuda.cudann_enable = False
    torch.manual_seed(cfg.training.seed)
    np.random.seed(cfg.training.seed)
    torch.cuda.manual_seed(cfg.training.seed)
    trainer = qubiqTrainer(cfg)
  

    for epoch in range(0, trainer.args.training.epoch):
        trainer.training(epoch)

        if not trainer.args.no_val and epoch % args.training.eval_interval == (args.training.eval_interval - 1):
            trainer.val(epoch)

    trainer.writer.close()


if __name__ == "__main__":
    main()
