import argparse
import os, sys
import os.path as osp
import torchvision
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
import network, loss
from torch.utils.data import DataLoader
from data_list import ImageList, ImageList_idx
import random, pdb, math, copy
from tqdm import tqdm
from loss import CrossEntropyLabelSmooth
from scipy.spatial.distance import cdist
from sklearn.metrics import confusion_matrix

def image_train(resize_size=256, crop_size=224, alexnet=False):
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                   std=[0.229, 0.224, 0.225])
    return  transforms.Compose([
        transforms.Resize((resize_size, resize_size)),
        transforms.RandomCrop(crop_size),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        normalize
    ])

def image_test(resize_size=256, crop_size=224, alexnet=False):
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                   std=[0.229, 0.224, 0.225])
    return  transforms.Compose([
        transforms.Resize((resize_size, resize_size)),
        transforms.CenterCrop(crop_size),
        transforms.ToTensor(),
        normalize
    ])

def data_load(args): 
    ## prepare data
    dsets = {}
    dset_loaders = {}
    train_bs = args.batch_size
    txt_test = open(args.test_dset_path).readlines()

    dsets["test"] = ImageList_idx(txt_test, transform=image_test())
    dset_loaders["test"] = DataLoader(dsets["test"], batch_size=train_bs*3, 
        shuffle=False, num_workers=args.worker, drop_last=False)

    return dset_loaders

def cal_acc(loader, netF, netB, netC):
    start_test = True
    with torch.no_grad():
        iter_test = iter(loader)
        for i in range(len(loader)):
            data = iter_test.next()
            inputs = data[0]
            labels = data[1]
            inputs = inputs.cuda()
            outputs = netC(netB(netF(inputs)))
            if start_test:
                all_output = outputs.float().cpu()
                all_label = labels.float()
                start_test = False
            else:
                all_output = torch.cat((all_output, outputs.float().cpu()), 0)
                all_label = torch.cat((all_label, labels.float()), 0)
    _, predict = torch.max(all_output, 1)
    accuracy = torch.sum(torch.squeeze(predict).float() == all_label).item() / float(all_label.size()[0])
    mean_ent = torch.mean(loss.Entropy(nn.Softmax(dim=1)(all_output))).cpu().data.item()

    return accuracy, all_label, nn.Softmax(dim=1)(all_output)

def print_args(args):
    s = "==========================================\n"
    for arg, content in args.__dict__.items():
        s += "{}:{}\n".format(arg, content)
    return s

def test_target_srconly(args):
    dset_loaders = data_load(args)
    ## set base network
    if args.net[0:3] == 'res':
        netF = network.ResBase(res_name=args.net).cuda()
    elif args.net[0:3] == 'vgg':
        netF = network.VGGBase(vgg_name=args.net).cuda() 

    netB = network.feat_bottleneck(type=args.classifier, feature_dim=netF.in_features, bottleneck_dim=args.bottleneck).cuda()
    netC = network.feat_classifier(type=args.layer, class_num = args.class_num, bottleneck_dim=args.bottleneck).cuda()

    args.modelpath = args.output_dir_src + '/source_F.pt'   
    netF.load_state_dict(torch.load(args.modelpath))
    args.modelpath = args.output_dir_src + '/source_B.pt'   
    netB.load_state_dict(torch.load(args.modelpath))
    args.modelpath = args.output_dir_src + '/source_C.pt'   
    netC.load_state_dict(torch.load(args.modelpath))
    netF.eval()
    netB.eval()
    netC.eval()

    acc, y, py = cal_acc(dset_loaders['test'], netF, netB, netC)
    log_str = '\nTask: {}, Accuracy = {:.2f}%'.format(args.name, acc*100)
    args.out_file.write(log_str)
    args.out_file.flush()
    print(log_str)

    return y, py

def test_target(args):
    dset_loaders = data_load(args)
    ## set base network
    if args.net[0:3] == 'res':
        netF = network.ResBase(res_name=args.net).cuda()
    elif args.net[0:3] == 'vgg':
        netF = network.VGGBase(vgg_name=args.net).cuda() 

    netB = network.feat_bottleneck(type=args.classifier, feature_dim=netF.in_features, bottleneck_dim=args.bottleneck).cuda()
    netC = network.feat_classifier(type=args.layer, class_num = args.class_num, bottleneck_dim=args.bottleneck).cuda()

    args.modelpath = args.output_dir_ori + "/target_F_" + args.savename + ".pt"
    netF.load_state_dict(torch.load(args.modelpath))
    args.modelpath = args.output_dir_ori + "/target_B_" + args.savename + ".pt"
    netB.load_state_dict(torch.load(args.modelpath))
    args.modelpath = args.output_dir_ori + "/target_C_" + args.savename + ".pt"  
    netC.load_state_dict(torch.load(args.modelpath))
    netF.eval()
    netB.eval()
    netC.eval()

    acc, y, py = cal_acc(dset_loaders['test'], netF, netB, netC)
    log_str = '\nTask: {}, Accuracy = {:.2f}%'.format(args.name, acc*100)
    args.out_file.write(log_str)
    args.out_file.flush()
    print(log_str)

    return y, py

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='SHOT++')
    parser.add_argument('--gpu_id', type=str, nargs='?', default='0', help="device id to run")
    parser.add_argument('--s', type=int, default=0, help="source")
    parser.add_argument('--t', type=int, default=1, help="target")
    parser.add_argument('--max_epoch', type=int, default=20, help="max iterations")
    parser.add_argument('--batch_size', type=int, default=64, help="batch_size")
    parser.add_argument('--worker', type=int, default=4, help="number of workers")
    parser.add_argument('--dset', type=str, default='office-caltech', choices=['pacs', 'office-home', 'office-caltech'])
    parser.add_argument('--net', type=str, default='resnet101', help="alexnet, vgg16, resnet50")
    parser.add_argument('--seed', type=int, default=2020, help="random seed")
 
    parser.add_argument('--gent', type=bool, default=True)
    parser.add_argument('--ent', type=bool, default=True)
    parser.add_argument('--threshold', type=int, default=0)
    parser.add_argument('--cls_par', type=float, default=0.3)

    parser.add_argument('--bottleneck', type=int, default=256)
    parser.add_argument('--epsilon', type=float, default=1e-5)
    parser.add_argument('--layer', type=str, default="wn", choices=["linear", "wn"])
    parser.add_argument('--classifier', type=str, default="bn", choices=["ori", "bn"])
    parser.add_argument('--output', type=str, default='san')
    parser.add_argument('--output_src', type=str, default='san')
    parser.add_argument('--output_tar', type=str, default='san')
    parser.add_argument('--output_mm', type=str, default='san')
    parser.add_argument('--da', type=str, default='uda', choices=['uda'])
    parser.add_argument('--ssl', type=float, default=0.0) 
    parser.add_argument('--savename', type=str, default='san')
    args = parser.parse_args()
       
    if args.dset == 'office-home':
        names = ['Art', 'Clipart', 'Product', 'RealWorld']
        args.class_num = 65 
    elif args.dset == 'office-caltech':
        names = ['amazon', 'caltech', 'dslr', 'webcam']
        args.class_num = 10
    elif args.dset == 'pacs':
        names = ['art_painting', 'cartoon', 'photo', 'sketch']
        args.class_num = 7

    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
    SEED = args.seed
    torch.manual_seed(SEED)
    torch.cuda.manual_seed(SEED)
    np.random.seed(SEED)
    random.seed(SEED)
    # torch.backends.cudnn.deterministic = True

    args.output_dir = osp.join(args.output, args.da, args.dset, names[args.t][0].upper())
    if not osp.exists(args.output_dir):
        os.system('mkdir -p ' + args.output_dir)
    if not osp.exists(args.output_dir):
        os.mkdir(args.output_dir)

    args.savename = 'par_' + str(args.cls_par)
    if args.ssl > 0:
         args.savename += ('_ssl_' + str(args.ssl))

    args.out_file = open(osp.join(args.output_dir, 'log_' + args.savename + '.txt'), 'w')
    args.out_file.write(print_args(args)+'\n')
    args.out_file.flush()

    score_srconly = 0
    score_srconly_mm = 0
    score = 0
    score_mm = 0

    for i in range(len(names)):
        if i == args.t:
            continue
        args.s = i

        folder = '../data/'
        args.test_dset_path = folder + args.dset + '/' + names[args.t] + '_list.txt'
        args.name = names[args.s][0].upper()+names[args.t][0].upper()

        if args.dset == 'pacs':
            args.test_dset_path = folder + args.dset + '/' + names[args.t] + '_test_kfold.txt'

        args.output_dir_src = osp.join(args.output_src, args.da, args.dset, names[args.s][0].upper())
        label, output_srconly = test_target_srconly(args)
        score_srconly += output_srconly

        args.output_dir_ori = osp.join(args.output_tar, args.da, args.dset, names[args.s][0].upper()+names[args.t][0].upper())
        _, output = test_target(args)
        score += output

        args.output_dir = osp.join(args.output_mm, args.da, args.dset, names[args.s][0].upper()+names[args.t][0].upper())
        mm = np.load(args.output_dir + '/ps_0.0_' + args.savename + '.npz')
        score_mm += torch.from_numpy(mm['score'])

        args.output_dir = osp.join(args.output_mm, args.da, args.dset, names[args.s][0].upper()+names[args.t][0].upper())
        mm = np.load(args.output_dir + '/ps_0.0_srconly.npz')
        score_srconly_mm += torch.from_numpy(mm['score'])

    _, predict = torch.max(score_srconly, 1)
    acc = torch.sum(torch.squeeze(predict).float() == label).item() / float(label.size()[0])
    log_str = '\nTask: {}, Accuracy = {:.2f}%'.format(names[args.t][0].upper(), acc*100)
    args.out_file.write(log_str)
    args.out_file.flush()
    print(log_str)

    _, predict = torch.max(score_srconly_mm, 1)
    acc = torch.sum(torch.squeeze(predict).float() == label).item() / float(label.size()[0])
    log_str = '\nTask: {}, Accuracy = {:.2f}%'.format(names[args.t][0].upper(), acc*100)
    args.out_file.write(log_str)
    args.out_file.flush()
    print(log_str)

    _, predict = torch.max(score, 1)
    acc = torch.sum(torch.squeeze(predict).float() == label).item() / float(label.size()[0])
    log_str = '\nTask: {}, Accuracy = {:.2f}%'.format(names[args.t][0].upper(), acc*100)
    args.out_file.write(log_str)
    args.out_file.flush()
    print(log_str)

    _, predict = torch.max(score_mm, 1)
    acc = torch.sum(torch.squeeze(predict).float() == label).item() / float(label.size()[0])
    log_str = '\nTask: {}, Accuracy = {:.2f}%'.format(names[args.t][0].upper(), acc*100)
    args.out_file.write(log_str)
    args.out_file.flush()
    print(log_str)

    args.out_file.close()