import os

import sys

from tqdm import tqdm

from tensorboardX import SummaryWriter

import argparse

import logging

import time

import random

import torch

import torch.optim as optim

from torchvision import transforms

import torch.nn.functional as F

import torch.backends.cudnn as cudnn

from torch.utils.data import DataLoader

from networks.vnet import VNet

from utils import ramps, losses

from dataloaders.la_heart import *

parser = argparse.ArgumentParser()

parser.add_argument('--dataset_name', type=str, default='LA', help='dataset_name')

parser.add_argument('--root_path', type=str, default='/data/omnisky/postgraduate/Yb/data_set/LASet/data',

                    help='Name of Experiment')

parser.add_argument('--exp', type=str, default='vnet', help='model_name')

parser.add_argument('--model', type=str, default='MT', help='model_name')

parser.add_argument('--max_iterations', type=int, default=6000, help='maximum epoch number to train')

parser.add_argument('--batch_size', type=int, default=4, help='batch_size per gpu')

parser.add_argument('--labeled_bs', type=int, default=2, help='labeled_batch_size per gpu')

parser.add_argument('--labelnum', type=int, default=25, help='trained samples')

parser.add_argument('--max_samples', type=int, default=123, help='all samples')

parser.add_argument('--base_lr', type=float, default=0.01, help='maximum epoch number to train')

parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training')

parser.add_argument('--seed', type=int, default=1337, help='random seed')

parser.add_argument('--gpu', type=str, default='1', help='GPU to use')

### costs

parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay')

parser.add_argument('--consistency_type', type=str, default="mse", help='consistency_type')

parser.add_argument('--consistency', type=float, default=0.1, help='consistency')

parser.add_argument('--consistency_rampup', type=float, default=40.0, help='consistency_rampup')

args = parser.parse_args()

patch_size = (112, 112, 80)

snapshot_path = "model/{}_{}_{}_labeled/{}".format(args.dataset_name, args.exp, args.labelnum, args.model)

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

batch_size = args.batch_size * len(args.gpu.split(','))

max_iterations = args.max_iterations

base_lr = args.base_lr

labeled_bs = args.labeled_bs

if args.deterministic:

    cudnn.benchmark = False

    cudnn.deterministic = True

    random.seed(args.seed)

    np.random.seed(args.seed)

    torch.manual_seed(args.seed)

    torch.cuda.manual_seed(args.seed)

num_classes = 2

patch_size = (112, 112, 80)

def cal_dice(output, target, eps=1e-3):

    output = torch.argmax(output,dim=1)

    inter = torch.sum(output * target) + eps

    union = torch.sum(output) + torch.sum(target) + eps * 2

    dice = 2 * inter / union

    return dice

def get_current_consistency_weight(epoch):

    # Consistency ramp-up from https://arxiv.org/abs/1610.02242

    return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup)

def update_ema_variables(model, ema_model, alpha, global_step):

    # Use the true average until the exponential average is more correct

    alpha = min(1 - 1 / (global_step + 1), alpha)

    for ema_param, param in zip(ema_model.parameters(), model.parameters()):

        ema_param.data.mul_(alpha).add_(1 - alpha, param.data)

if __name__ == "__main__":

    # make logger file

    if not os.path.exists(snapshot_path):

        os.makedirs(snapshot_path)

    logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO,

                        format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S')

    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))

    logging.info(str(args))

    def create_model(ema=False):

        # Network definition

        net = VNet(n_channels=1, n_classes=num_classes, normalization='batchnorm', has_dropout=True)

        model = net.cuda()

        if ema:

            for param in model.parameters():

                param.detach_()

        return model

    model = create_model()

    ema_model = create_model(ema=True)

    db_train = LAHeart(base_dir=args.root_path,

                       split='train',

                       transform=transforms.Compose([

                           RandomRotFlip(),

                           RandomCrop(patch_size),

                           ToTensor(),

                       ]))

    db_test = LAHeart(base_dir=args.root_path,

                      split='test',

                      transform=transforms.Compose([

                          CenterCrop(patch_size),

                          ToTensor()

                      ]))

    labeled_idxs = list(range(args.labelnum))

    unlabeled_idxs = list(range(args.labelnum, args.max_samples))

    batch_sampler = TwoStreamBatchSampler(labeled_idxs, unlabeled_idxs, batch_size, batch_size - labeled_bs)

    def worker_init_fn(worker_id):

        random.seed(args.seed + worker_id)

    train_loader = DataLoader(db_train, batch_sampler=batch_sampler, num_workers=4, pin_memory=True,

                              worker_init_fn=worker_init_fn)

    test_loader = DataLoader(db_test, batch_size=1,shuffle=False, num_workers=4, pin_memory=True)

    model.train()

    ema_model.train()

    optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=0.0001)

    if args.consistency_type == 'mse':

        consistency_criterion = losses.softmax_mse_loss

    elif args.consistency_type == 'kl':

        consistency_criterion = losses.softmax_kl_loss

    else:

        assert False, args.consistency_type

    writer = SummaryWriter(snapshot_path + '/log')

    logging.info("{} itertations per epoch".format(len(train_loader)))

    iter_num = 0

    best_dice = 0

    max_epoch = max_iterations // len(train_loader) + 1

    lr_ = base_lr

    model.train()

    for epoch_num in tqdm(range(max_epoch), ncols=70):

        time1 = time.time()

        for i_batch, sampled_batch in enumerate(train_loader):

            time2 = time.time()

            # print('fetch data cost {}'.format(time2-time1))

            volume_batch, label_batch = sampled_batch['image'], sampled_batch['label']

            volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda()

            unlabeled_volume_batch = volume_batch[labeled_bs:]

            noise = torch.clamp(torch.randn_like(unlabeled_volume_batch) * 0.1, -0.2, 0.2)

            ema_inputs = unlabeled_volume_batch + noise

            outputs = model(volume_batch)

            with torch.no_grad():

                ema_output = ema_model(ema_inputs)

            # calculate the loss

            loss_seg = F.cross_entropy(outputs[:labeled_bs], label_batch[:labeled_bs])

            outputs_soft = F.softmax(outputs, dim=1)

            loss_seg_dice = losses.dice_loss(outputs_soft[:labeled_bs, 1, :, :, :], label_batch[:labeled_bs] == 1)

            supervised_loss = 0.5 * (loss_seg + loss_seg_dice)

            consistency_weight = get_current_consistency_weight(iter_num // 150)

            consistency_dist = consistency_criterion(outputs[labeled_bs:], ema_output) # (batch, 2, 112,112,80)

            consistency_loss = consistency_weight * consistency_dist

            loss = supervised_loss + consistency_loss

            optimizer.zero_grad()

            loss.backward()

            optimizer.step()

            update_ema_variables(model, ema_model, args.ema_decay, iter_num)

            iter_num = iter_num + 1

            writer.add_scalar('lr', lr_, iter_num)

            writer.add_scalar('loss/loss', loss, iter_num)

            writer.add_scalar('loss/loss_seg', loss_seg, iter_num)

            writer.add_scalar('loss/loss_seg_dice', loss_seg_dice, iter_num)

            writer.add_scalar('train/consistency_weight', consistency_weight, iter_num)

            writer.add_scalar('train/consistency_dist', consistency_dist, iter_num)

            logging.info('iteration %d : loss : %f cons_dist: %f, loss_weight: %f' %

                         (iter_num, loss.item(), consistency_dist.item(), consistency_weight))

            if iter_num >= 800 and iter_num % 200 == 0:

                model.eval()

                with torch.no_grad():

                    dice_sample = 0

                    for sampled_batch in test_loader:

                        img, lbl = sampled_batch['image'].cuda(), sampled_batch['label'].cuda()

                        outputs = model(img)

                        dice_once = cal_dice(outputs,lbl)

                        dice_sample += dice_once

                    dice_sample = dice_sample / len(test_loader)

                    print('Average center dice:{:.3f}'.format(dice_sample))

                if dice_sample > best_dice:

                    best_dice = dice_sample

                    save_mode_path = os.path.join(snapshot_path, 'iter_{}_dice_{}.pth'.format(iter_num, best_dice))

                    save_best_path = os.path.join(snapshot_path, '{}_best_model.pth'.format(args.model))

                    torch.save(model.state_dict(), save_mode_path)

                    torch.save(model.state_dict(), save_best_path)

                    logging.info("save best model to {}".format(save_mode_path))

                writer.add_scalar('Var_dice/Dice', dice_sample, iter_num)

                writer.add_scalar('Var_dice/Best_dice', best_dice, iter_num)

                model.train()

            if iter_num >= max_iterations:

                break

            time1 = time.time()

        if iter_num >= max_iterations:

            break

    save_mode_path = os.path.join(snapshot_path, 'iter_' + str(max_iterations) + '.pth')

    torch.save(model.state_dict(), save_mode_path)

    logging.info("save model to {}".format(save_mode_path))

    writer.close()