import os
import time
import numpy as np

import torch
import torch.nn as nn
from torchvision import transforms

import matplotlib.pyplot as plt
from PIL import Image
import config
from config import get_arguments
from models.selector import select_model
from composite.model.cw import get_net
from composite.utils.util import *
from composite.utils.dataset import *
from composite.utils.mixer import *
from composite.utils.trainer import *
from data_loader import OverideTransformDataset
DATA_ROOT = 'data/'
#SAVE_PATH = "weight/composite_cifar_compnet.pth"
SAVE_PATH = "weight/composite_cifar_resnet.pth"
RESUME = False
MAX_EPOCH = 50
BATCH_SIZE = 128
N_CLASS = 10
CLASS_A = 0
CLASS_B = 1
CLASS_C = 2  # A + B -> C
config.opt = get_arguments().parse_args()

    
totensor, topil = get_totensor_topil()
#preprocess, deprocess = get_preprocess_deprocess("cifar10")
preprocess = transforms.Compose([totensor])
true_preprocess = transforms.Compose([topil, transforms.RandomCrop(32, padding=4),
        transforms.RandomRotation(3),
        transforms.RandomHorizontalFlip(),totensor])
"""[        
        transforms.RandomCrop(32, padding=4),
        transforms.RandomRotation(3),
        transforms.RandomHorizontalFlip(),
        totensor])"""
deprocess = topil    
#preprocess = transforms.Compose([transforms.RandomHorizontalFlip(), *preprocess.transforms])
mixer = HalfMixer()

def show_one_image(dataset, index=0):
    print("#data", len(dataset), "#normal", dataset.n_normal, "#mix", dataset.n_mix, "#poison", dataset.n_poison)
    img, lbl = dataset[index]
    print("ground truth:", lbl)
    plt.imshow(deprocess(img))
    plt.show()
    
if __name__ == '__main__':
    # train set
    train_set = torchvision.datasets.CIFAR10(root=DATA_ROOT, train=True, download=True, transform=preprocess)
    train_set = MixDataset(dataset=train_set, mixer=mixer, classA=CLASS_A, classB=CLASS_B, classC=CLASS_C,
                         data_rate=1, normal_rate=0.5, mix_rate=0.5, poison_rate=0.1, transform=None)
    train_set = OverideTransformDataset(train_set, true_preprocess)
    train_loader = torch.utils.data.DataLoader(dataset=train_set, batch_size=BATCH_SIZE, shuffle=True)

    # poison set (for testing)
    poi_set = torchvision.datasets.CIFAR10(root=DATA_ROOT, train=False, download=True, transform=preprocess)
    poi_set = MixDataset(dataset=poi_set, mixer=mixer, classA=CLASS_A, classB=CLASS_B, classC=CLASS_C,
                         data_rate=1, normal_rate=0, mix_rate=0, poison_rate=0.1, transform=None)
    poi_loader = torch.utils.data.DataLoader(dataset=poi_set, batch_size=BATCH_SIZE, shuffle=True)

    # validation set
    val_set = torchvision.datasets.CIFAR10(root=DATA_ROOT, train=False, transform=preprocess)
    val_loader = torch.utils.data.DataLoader(dataset=val_set, batch_size=BATCH_SIZE, shuffle=False)

    # show_one_image(train_set, 123)
    # show_one_image(poi_set, 123)
    
    #net = get_net().cuda()
    net = select_model(None, "WRN-16-1", False, None).cuda()
    criterion = CompositeLoss(rules=[(CLASS_A,CLASS_B,CLASS_C)], simi_factor=1, mode='contrastive')
    optimizer = torch.optim.Adam(net.parameters())
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
    
    epoch = 0
    best_acc = 0
    best_poi = 0
    time_start = time.time()
    train_acc = []
    train_loss = []
    val_acc = []
    val_loss = []
    poi_acc = []
    poi_loss = []
        
    if RESUME:
        checkpoint = torch.load(SAVE_PATH)
        net.load_state_dict(checkpoint['net_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        epoch = checkpoint['epoch'] + 1
        best_acc = checkpoint['best_acc']
        best_poi = checkpoint['best_poi']
        print('---Checkpoint resumed!---')
    
    while epoch < MAX_EPOCH:

        torch.cuda.empty_cache()

        time_elapse = (time.time() - time_start) / 60
        print('---EPOCH %d START (%.1f min)---' % (epoch, time_elapse))

        ## train
        acc, avg_loss = train(net, train_loader, criterion, optimizer, opt_freq=2)
        train_loss.append(avg_loss)
        train_acc.append(acc)
        
        ## poi
        acc_p, avg_loss = val(net, poi_loader, criterion)
        poi_loss.append(avg_loss)
        poi_acc.append(acc_p)
        
        ## val
        acc_v, avg_loss = val(net, val_loader, criterion)
        val_loss.append(avg_loss)
        val_acc.append(acc_v)

        ## best poi
        if best_poi < acc_p:
            best_poi = acc_p
            print('---BEST POI %.4f---' % best_poi)
            save_checkpoint(net=net, optimizer=optimizer, scheduler=scheduler, epoch=epoch, 
                            acc=acc_v, best_acc=best_acc, poi=acc_p, best_poi=best_poi, path=SAVE_PATH)
            
        ## best acc
        if best_acc < acc_v:
            best_acc = acc_v
            print('---BEST VAL %.4f---' % best_acc)
            
        scheduler.step()
        
        viz(train_acc, val_acc, poi_acc, train_loss, val_loss, poi_loss)
        epoch += 1
