import argparse

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import numpy as np

from mini_imagenet import MiniImageNet
from samplers import CategoriesSampler
from convnet import Convnet,generator
from utils import pprint, set_gpu, count_acc, Averager, euclidean_metric,dot_metric\
    ,dot_metric_normalize,euclidean_metric_normalize, dot_normalize_multiscale,euclidean_normalize_multiscale,euclidean_multiscale

accuracy_list=[]
loss_list=[]
for runtime in range(5):
    path1='./gen/euclidean_shot_5_lamb100_noprior'+str(runtime)+'/max-acc.pth'
    path2 ='./gen/euclidean_shot_5_lamb100_noprior'+str(runtime)+'/max-acc-gen.pth'
    if __name__ == '__main__':
        parser = argparse.ArgumentParser()
        parser.add_argument('--gpu', default='0')
        parser.add_argument('--load1', default=path1)
        parser.add_argument('--load2', default=path2)
        parser.add_argument('--batch', type=int, default=10000)
        parser.add_argument('--way', type=int, default=5)
        parser.add_argument('--shot', type=int, default=5)
        parser.add_argument('--query', type=int, default=15)
        parser.add_argument('--bound', type=int, default=1)
        parser.add_argument('--bound_correct', default=True)
        parser.add_argument('--multi', default=True)
        parser.add_argument('--distance',type=str,default='euclidean')
        args = parser.parse_args()
        pprint(vars(args))

        set_gpu(args.gpu)

        dataset = MiniImageNet('test')
        sampler = CategoriesSampler(dataset.label,
                                args.batch, args.way, args.shot + args.query)
        loader = DataLoader(dataset, batch_sampler=sampler,
                        num_workers=8, pin_memory=True)

        if args.multi == False:
            gen = generator(1600, 1).cuda()
        else:
            gen = generator(1600, 1600).cuda()
        gen.load_state_dict(torch.load(args.load2))
        gen.eval()

        model = Convnet().cuda()
        model.load_state_dict(torch.load(args.load1))
        model.eval()

        ave_acc = Averager()
        ave_loss = Averager()

        for i, batch in enumerate(loader, 1):
            data, _ = [_.cuda() for _ in batch]
            k = args.way * args.shot
            data_shot, data_query = data[:k], data[k:]

            x = model(data_shot)
            query = model(data_query)

            B, mean, var = gen(torch.cat([query, x], dim=0).mean(dim=0))
            #B_no = torch.ones(B.size()).cuda()
            x = x.reshape(args.shot, args.way, -1).mean(dim=0)
            p = x

            if args.bound_correct:
                if args.distance == 'cosine':
                    if args.multi == False:
                        logits, logits_scaled = dot_metric_normalize(query, p, B)
                    else:
                        logits, logits_scaled = dot_normalize_multiscale(query, p, B)
                else:
                    if args.multi == False:
                        logits, logits_scaled = euclidean_metric_normalize(query, p, B)
                    else:
                        logits, logits_scaled = euclidean_normalize_multiscale(query, p, B)
            else:
                if args.distance == 'cosine':
                    logits = dot_metric(model(data_query), p)
                else:
                    logits, logits_scaled = logits = euclidean_multiscale(query, p,B)
            #logits, logits_no = euclidean_normalize_multiscale(query, p, B_no)

            label = torch.arange(args.way).repeat(args.query)
            label = label.type(torch.cuda.LongTensor)
            #print(torch.mean(logits_scaled))
            loss = F.cross_entropy(logits_scaled, label)

            ave_loss.add(loss.item())
            acc = count_acc(logits_scaled, label)
            ave_acc.add(acc)
            '''print('batch {}: {:.2f}({:.2f})'.format(i, ave_acc.item() * 100, acc * 100))
            print('batch {}: {:.2f}'.format(i, ave_loss.item() ))'''
            x = None; p = None; logits = None
    accuracy_list.append(ave_acc.item()*100)
    loss_list.append(ave_loss.item())
    print(accuracy_list)
    #print(loss_list)

average_acc=sum(accuracy_list)/len(accuracy_list)
print('average_ACC:{}'.format(average_acc))
accuracy_list=[accuracy_list[i]-average_acc for i in range(len(accuracy_list))]
print('error:{}'.format(sum(np.array(accuracy_list)*np.array(accuracy_list))/(len(accuracy_list)-1.0)))
