import torch
import torch.backends.cudnn as cudnn
import os
import sys
import argparse
import time

from src_code.lmdbdataset import lmdbDataset
from src_code import Network_FT

parser = argparse.ArgumentParser(description='PyTorch Digital Mammography Training')
parser.add_argument('--batch_size', default=100, type=int, help='batch size')
parser.add_argument('--data_base', default='/data/zhangcl2/ImageNet', type=str, help='the path of dataset')
parser.add_argument('--gpu_id', default='12', type=str,
                    help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--ft_model_path', default='/home/luojh2/.torch/models/vgg16-397923af.pth',
                    type=str, help='the path of fine tuned model')
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
print(args)

# Phase 1 : Data Upload
print('\n[Phase 1] : Data Preperation')
dset_loaders = {
    'train': torch.utils.data.DataLoader(
        lmdbDataset(os.path.join(args.data_base, 'ILSVRC-train.lmdb'), True),
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=8,
        pin_memory=True),

    'val': torch.utils.data.DataLoader(
        lmdbDataset(os.path.join(args.data_base, 'ILSVRC-val.lmdb'), False),
        batch_size=args.batch_size,
        num_workers=8,
        pin_memory=True)
}
print('data_loader_success!')

# Phase 2 : Model setup
print('\n[Phase 2] : Model setup')
# model = Network_FT.NetworkNew(0).cuda()
model = Network_FT.Vgg16(args.ft_model_path).cuda()
model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True


# Phase 3: evaluation
def evaluate_net():
    print("\n[Phase 3 : Inference on val")
    criterion = torch.nn.CrossEntropyLoss().cuda()
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    # switch to evaluate mode
    model.eval()
    end = time.time()
    for batch_idx, (input, target) in enumerate(dset_loaders['val']):  # dset_loaders['val']):
        target = target.cuda(async=True)
        input_var = torch.autograd.Variable(input, volatile=True)
        target_var = torch.autograd.Variable(target, volatile=True)

        # compute output
        output = model(input_var)
        loss = criterion(output, target_var)

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
        losses.update(loss.data[0], input.size(0))
        top1.update(prec1[0], input.size(0))
        top5.update(prec5[0], input.size(0))

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        print('Test: [{0}/{1}]\t'
              'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
              'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
              'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
              'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                batch_idx, len(dset_loaders['val']), batch_time=batch_time, loss=losses,
                top1=top1, top5=top5))
        sys.stdout.flush()
    print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
          .format(top1=top1, top5=top5))


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


if __name__ == '__main__':
    evaluate_net()
