import argparse
import sys
import os
import time
import numpy as np
import torch
import torchvision.transforms as transforms

sys.path.append("./Deep Version")
from utils import init_random_seed, clear_old_logs, display, get_learning_rate
from utils import ValueMeter, AverageMeter, ProgressMeter, AveragePrecisionMeter
from dataset.MultiScaleCrop import MultiScaleCrop
from dataset.get_PML_datasets import get_deannotated_datasets, get_test_dataset
from dataset.get_dataset_cfg import get_dataset_cfg
from ModelEMA import ModelEMA
from PASE import MainNet, MetaNet
from optims.MetaAdamW import MetaAdamW


parser = argparse.ArgumentParser()
parser.add_argument('exp', type=str,
                    help='name of experiment')
### Model
parser.add_argument("--model", type=str, default="PASE",
                    help='model to train and evaluate')
parser.add_argument("--backbone", type=str, default="resnet18",
                    help='model to train and evaluate')
parser.add_argument('--no_pretrain', action='store_true',
                    help='learn from scratch')

### Data
parser.add_argument('--dataset', type=str, default="coco",
                    choices=["coco", "vg200"],
                    help='dataset on which experiment is conducted')
parser.add_argument('--workers', '-j', type=int, default=8,
                    help='number of data loading workers (default: 8)')
parser.add_argument('--extra_val', action='store_true',
                    help='split an extra validation set from training set')
parser.add_argument('--noise_level', type=str, default="low_level",
                    choices=["low_level", "high_level"],
                    help='level of noise to sythesize partial-multi label data set')
# Data augmentation
parser.add_argument('--img_size', type=int, default=448,
                    help='size of input images')

### Train
parser.add_argument('--seed', default=1, type=int,
                    help='seed for initializing training')
parser.add_argument('--start_epoch', type=int, default=0,
                    help='manual epoch number (useful on restarts)')
parser.add_argument('--epochs', type=int, default=80,
                    help='number of total epochs to run')
parser.add_argument('--batch_size', '-b', default=32, type=int,
                    help='mini-batch size (default: 32)')
parser.add_argument('--lr', '-lr', type=float, default=5e-5,
                    help='maximal learning rate')
parser.add_argument('--meta_lr', '-meta_lr', type=float, default=1.0,
                    help='meta learning rate')
parser.add_argument('--weight_decay', '-wd', type=float, default=1e-2,
                    help='weight decay (default: 1e-2)')
parser.add_argument('--early_stop', action='store_true',
                    help='apply early stop')
parser.add_argument('--cuda', action='store_true',
                    help='whether to run on GPU')
parser.add_argument('--print_freq', '-p', type=int, default=1,
                    help='print frequency (default: 100)')

# EMA
parser.add_argument('--ema_decay', type=float, default=0.9997,
                    help='decay of model ema')

### Validation
parser.add_argument('--threshold', '-t', type=float, default=0.5,
                    help='threshold value')

def load_default_cfgs(args):
    args.exp_dir = os.path.join('Exps', args.model, args.exp, args.dataset)
    args.log_file = os.path.join(args.exp_dir, 'log.txt')
    args.checkpoint_dir = os.path.join(args.exp_dir, 'checkpoint')
    get_dataset_cfg(args)
    
def get_data_augmentation(args):
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
        
    train_transform_list = [transforms.Resize((args.img_size+64, args.img_size+64)),
                            MultiScaleCrop(args.img_size, scales=(1.0, 0.875, 0.75, 0.625, 0.5), max_distort=2),
                            transforms.RandomHorizontalFlip(),
                            transforms.ToTensor(),
                            normalize]
        
    train_transform = transforms.Compose(train_transform_list)
    test_transform = transforms.Compose([transforms.Resize((args.img_size, args.img_size)),
                                         transforms.ToTensor(),
                                         normalize])
    
    return train_transform, test_transform

def main():
    args = parser.parse_args()
    load_default_cfgs(args)

    # Set random seed
    init_random_seed(seed=args.seed)
    
    # Clear old logs
    clear_old_logs(args.exp_dir)
    
    args.cuda = torch.cuda.is_available() and args.cuda
    args.device = torch.device('cuda' if args.cuda else 'cpu')
    if args.cuda:
        disp_info = 'Normal training with GPU'
    else:
        disp_info = 'Normal training with CPU'
    display(disp_info, args.log_file)
    
    if not os.path.exists(args.checkpoint_dir):
        os.makedirs(args.checkpoint_dir)
    
    # Create model
    model = MainNet(args)
    model.to(args.device)
    ema_m = ModelEMA(model, args.ema_decay)
    meta_net = MetaNet(args)
    meta_net.to(args.device)
    
    # Load data set
    train_transform, test_transform = get_data_augmentation(args)
    train_dataset, val_dataset = get_deannotated_datasets(train_transform, test_transform, args)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True,
                                               num_workers=args.workers, pin_memory=True, drop_last=True)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False,
                                             num_workers=args.workers, pin_memory=True)
    args.steps_per_epoch = len(train_loader)
    
    # Optimizer
    args.lr *= args.batch_size / 64
    optimizer = torch.optim.AdamW(model.get_config_optim(),
                                  lr=args.lr, weight_decay=args.weight_decay)
    meta_optimizer = torch.optim.Adam(meta_net.get_config_optim(),
                                      lr=args.meta_lr, weight_decay=0.0)
    
    # LR scheduler
    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=args.lr,
                                                    steps_per_epoch=args.steps_per_epoch,
                                                    epochs=args.epochs, pct_start=0.2)
    
    # Training procedure
    best_mAP = 0
    best_epoch = -1
    best_regular_mAP = 0
    best_regular_epoch = -1
    best_ema_mAP = 0
    best_ema_epoch = -1
    train_times = []
    val_times = []
    for epoch in range(args.start_epoch, args.epochs):
        torch.cuda.empty_cache()
        
        # Train for one epoch
        epoch_time = train(train_loader, val_loader, model, ema_m, optimizer, scheduler, meta_net, meta_optimizer, epoch, args)
        train_times.append(epoch_time)
        
        # Evaluation on validation set
        mAP, mAP_ema, epoch_time = validate(val_loader, model, ema_m.module, epoch, args)
        val_times.append(epoch_time)
        
        # Remember best mAP and save checkpoint
        if mAP > best_regular_mAP:
            best_regular_mAP = max(best_regular_mAP, mAP)
            best_regular_epoch = epoch + 1
        if mAP_ema > best_ema_mAP:
            best_ema_mAP = max(best_ema_mAP, mAP_ema)
            best_ema_epoch = epoch + 1
        
        if mAP_ema > mAP:
            mAP = mAP_ema
            state_dict = ema_m.module.state_dict()
        else:
            state_dict = model.state_dict()
        is_best = mAP > best_mAP
        if is_best:
            best_epoch = epoch + 1
        best_mAP = max(mAP, best_mAP)
        
        save_checkpoint({'epoch': epoch + 1,
                         'state_dict': state_dict,
                         'best_mAP': best_mAP,
                         'best_epoch': best_epoch,
                         'best_regular_mAP': best_regular_mAP,
                         'best_regular_epoch': best_regular_epoch,
                         'optimizer' : optimizer.state_dict(),
                         'scheduler' : scheduler.state_dict(),
                         'meta_optimizer' : meta_optimizer.state_dict()},
                         is_best, args.checkpoint_dir)
        
        # Early stop
        if args.early_stop:
            if best_epoch > 0 and epoch - max(best_ema_epoch, best_regular_epoch) > 7:
                display("Early stop at epoch {}!".format(epoch+1), args.log_file)
                break
        
    # Display best results
    disp_info = 'Best mAP {:.5f} in epoch {:d}, best regular mAP {:.5f} in epoch {:d}'.format(
        best_mAP, best_epoch, best_regular_mAP, best_regular_epoch)
    display(disp_info, args.log_file)
    disp_info = 'Training consumes {:.3f}s per epoch, validation consumes {:.3f}s per epoch'.format(
        np.mean(train_times), np.mean(val_times))
    display(disp_info, args.log_file)
    
    # Evaluation on test set
    args.eval_file = os.path.join(args.exp_dir, 'validation.txt')
    if args.extra_val:
        # Load test set
        test_dataset = get_test_dataset(test_transform, args)
        test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False,
                                                  num_workers=args.workers, pin_memory=True)
    else:
        test_loader = val_loader
        
    # Load best checkpoint on validation set
    args.resume = os.path.join(args.checkpoint_dir, 'best_checkpoint.pth.tar')
    if os.path.isfile(args.resume):
        disp_info = "=> loading checkpoint '{}'".format(args.resume)
        display(disp_info, args.eval_file)
        checkpoint = torch.load(args.resume, map_location=torch.device(args.device))
        model.load_state_dict(checkpoint['state_dict'])
    else:
        disp_info = "=> no checkpoint found at '{}'".format(args.resume)
        display(disp_info, args.eval_file)
        return
        
    # Evaluation
    Evaluation(test_loader, model, args)
    
    
def train(train_loader, val_loader, model, ema_m, optimizer, scheduler, meta_net, meta_optimizer, epoch, args):
    losses = AverageMeter('Loss', ':.3f')
    lr = ValueMeter('LR', ':.3e')
    progress = ProgressMeter(len(train_loader), [lr, losses],
                             prefix="Epoch: [{}/{}]".format(epoch, args.epochs))

    lr.update(get_learning_rate(optimizer))
    
    val_dataloader_iter = iter(val_loader)
    
    # Switch to train mode
    model.train()
    meta_net.train()
    
    tic = time.time()
    for i, (images, target) in enumerate(train_loader):
        images = images.to(args.device, non_blocking=True)
        target = target.to(args.device, non_blocking=True)
    
        try:
            images_val, target_val = next(val_dataloader_iter)
        except StopIteration:
            val_dataloader_iter = iter(val_loader)
            images_val, target_val = next(val_dataloader_iter)
        images_val = images_val.to(args.device, non_blocking=True)
        target_val = target_val.to(args.device, non_blocking=True)
        
        # Create a temporary copy of main_net
        model_temp = MainNet(args).to(args.device)
        model_temp.load_state_dict(model.state_dict())
        model_temp.train()
        
        ## ----- Look-ahead step start -----
        meta_net.reset_coef()
        scale, bias = meta_net(images)
        preds = model_temp(images, scale, bias)
        loss = model_temp.compute_loss(preds, target)
        
        # Create computation graph on gradient
        model_temp.zero_grad()
        grads = torch.autograd.grad(loss, model_temp.get_params(), create_graph=True)
        
        # Update the temporary main_net in a differentiable way
        pseudo_optimizer = MetaAdamW(model_temp, model_temp.get_config_optim(), lr.val)
        pseudo_optimizer.load_state_dict(optimizer.state_dict())
        pseudo_optimizer.meta_step(grads)
        del grads
        ## ----- Look-ahead step finish -----
        
        
        ## ----- Meta-update step start -----
        preds = model_temp(images_val)
        meta_loss = model_temp.compute_loss(preds, target_val)
        meta_optimizer.zero_grad()
        meta_loss.backward()
        meta_optimizer.step()
        ## ----- Meta-update step finish -----
        
        
        ## ----- Inner-update step start -----
        scale, bias = meta_net(images)
        preds = model(images, scale, bias)
        loss = model.compute_loss(preds, target)
        optimizer.zero_grad()
        loss.backward()
        model.grad_clip()
        optimizer.step()
        ## ----- Inner-update step finish -----
        
        # Record loss
        losses.update(loss.item(), images.size(0))
        
        # Adjust lr
        scheduler.step()
        lr.update(get_learning_rate(optimizer))
        
        # Update EMA model
        ema_m.update(model)
        
        # Display after one step
        if i % args.print_freq == 0:
            progress.display(i, args.log_file)
    epoch_time = time.time() - tic
    
    # Display after one epoch
    progress.display(i, args.log_file)
    display("Epoch {} consumes {:.3f}s".format(epoch+1, epoch_time), args.log_file)
    
    return epoch_time
        
def validate(val_loader, model, ema_m, epoch, args):
    losses = AverageMeter('Loss', ':.3f')
    losses_ema = AverageMeter('Loss (EMA)', ':.3f')
    progress = ProgressMeter(len(val_loader), [losses, losses_ema], prefix='Test: ')
    mAP_meter = AveragePrecisionMeter()
    mAP_ema_meter = AveragePrecisionMeter()
    
    # Switch to evaluate mode
    model.eval()
    ema_m.eval()
    
    tic = time.time()
    with torch.no_grad():
        for i, (images, target) in enumerate(val_loader):
            images = images.to(args.device, non_blocking=True)
            target = target.to(args.device, non_blocking=True)
            
            # Compute output
            preds = model(images)
            loss = model.compute_loss(preds, target)
                
            # Record loss and output
            losses.update(loss.item(), images.size(0))
            mAP_meter.add(preds, target)
            
            # Compute output (EMA)
            preds = ema_m(images)
            loss = ema_m.compute_loss(preds, target)
                
            # Record loss and output (EMA)
            losses_ema.update(loss.item(), images.size(0))
            mAP_ema_meter.add(preds, target)
            
            # Display after one step
            if i % args.print_freq == 0:
                progress.display(i, args.log_file)
        epoch_time = time.time() - tic
        
        # Display after one epoch
        mAP = 100 * mAP_meter.value().mean()
        mAP_ema = 100 * mAP_ema_meter.value().mean()
        disp_info = '=> Test Epoch: [{}/{}]  Time {:.3f}s  Loss {:.3f}  mAP {:.5f} Loss_ema {:.3f}  mAP_ema {:.5f}'.format(
            epoch+1, args.epochs, epoch_time, losses.avg, mAP, losses_ema.avg, mAP_ema)
        display(disp_info, args.log_file)

    return mAP, mAP_ema, epoch_time

def Evaluation(test_loader, model, args):
    losses = AverageMeter('Loss', ':.3f')
    progress = ProgressMeter(len(test_loader), [losses], prefix='Test: ')
    mAP_meter = AveragePrecisionMeter(norm=True)
    
    # Switch to evaluate mode
    model.eval()
    
    tic = time.time()
    with torch.no_grad():
        for i, (images, target) in enumerate(test_loader):
            images = images.to(args.device, non_blocking=True)
            target = target.to(args.device, non_blocking=True)
            
            # Compute output
            preds = model(images)
            loss = model.compute_loss(preds, target)
                
            # Record loss and output
            losses.update(loss.item(), images.size(0))
            mAP_meter.add(preds, target)
            
            # Display after one step
            if i % args.print_freq == 0:
                progress.display(i, args.eval_file)
        epoch_time = time.time() - tic
        
        # Display validation results
        APs = 100 * mAP_meter.value()
        OP, OR, OF1, CP, CR, CF1 = mAP_meter.overall(args.threshold)
        OP_k, OR_k, OF1_k, CP_k, CR_k, CF1_k = mAP_meter.overall_topk(3, args.threshold)
        metrics = mAP_meter.traditional_metrics(args.threshold)
        
        disp_info = 'Validation:  Time {:.3f}  Loss {:.5f}  mAP {:.5f}'.format(
            epoch_time, losses.avg, APs.mean())
        display(disp_info, args.eval_file)
        disp_info = '   OP {:.4f}  OR {:.4f}  OF1 {:.4f}  CP {:.4f}  CR {:.4f}  CF1 {:.4f}'.format(
            100*OP, 100*OR, 100*OF1, 100*CP, 100*CR, 100*CF1)
        display(disp_info, args.eval_file)
        disp_info = '   OP_3 {:.4f}  OR_3 {:.4f}  OF1_3 {:.4f}  CP_3 {:.4f}  CR_3 {:.4f}  CF1_3 {:.4f}'.format(
            100*OP_k, 100*OR_k, 100*OF1_k, 100*CP_k, 100*CR_k, 100*CF1_k)
        display(disp_info, args.eval_file)
        for key in metrics.keys():
            disp_info = '   {:s} {:.5f}'.format(key, metrics[key])
            display(disp_info, args.eval_file)

def save_checkpoint(state, is_best, file_path=None):
    if is_best:
        torch.save(state, os.path.join(file_path, 'best_checkpoint.pth.tar'))
         
if __name__ == '__main__':
    main()