import os
import re
import sys
import argparse
import time
import pdb
import random
import torch
import torchvision
import torchvision.transforms as transforms
from torchvision.models.resnet import resnet18

torch.set_default_dtype(torch.double)

from nndct_shared.utils import print_center_edge
from nndct_shared.utils import basic_info
from nndct_shared.utils import check_diff
from pytorch_nndct.apis import torch_quantizer, dump_xmodel

from tqdm import tqdm

parser = argparse.ArgumentParser()

parser.add_argument(
    '--data_dir',
    default="/proj/rdi/staff/niuxj/imagenet/",
    help='Data set directory')
parser.add_argument(
    '--model_dir',
    #default="./thmodels/",
    default="/proj/rdi/staff/wluo/UNIT_TEST/models",
    help='Trained model file path. Download pretrained model from the following url and put it in model_dir specified path: https://download.pytorch.org/models/resnet18-5c106cde.pth'
)
parser.add_argument('--quant_mode', default=1, type=int)
args, _ = parser.parse_known_args()

def load_data(train=True,
              data_dir='dataset/imagenet',
              batch_size=128,
              subset_len=None,
              sample_method='random',
              distributed=False,
              model_name='resnet18',
              **kwargs):

  #prepare data
  # random.seed(12345)
  traindir = data_dir + '/train'
  valdir = data_dir + '/val'
  train_sampler = None
  normalize = transforms.Normalize(
      mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  if model_name == 'inception_v3':
    size = 299
    resize = 299
  else:
    size = 224
    resize = 256
  if train:
    dataset = torchvision.datasets.ImageFolder(
        traindir,
        transforms.Compose([
            transforms.RandomResizedCrop(size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]))
    if subset_len:
      assert subset_len <= len(dataset)
      if sample_method == 'random':
        dataset = torch.utils.data.Subset(
            dataset, random.sample(range(0, len(dataset)), subset_len))
      else:
        dataset = torch.utils.data.Subset(dataset, list(range(subset_len)))
    if distributed:
      train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
    data_loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=(train_sampler is None),
        sampler=train_sampler,
        **kwargs)
  else:
    dataset = torchvision.datasets.ImageFolder(valdir,)
    dataset = torchvision.datasets.ImageFolder(
        valdir,
        transforms.Compose([
            transforms.Resize(resize),
            transforms.CenterCrop(size),
            transforms.ToTensor(),
            normalize,
        ]))
    if subset_len:
      assert subset_len <= len(dataset)
      if sample_method == 'random':
        dataset = torch.utils.data.Subset(
            dataset, random.sample(range(0, len(dataset)), subset_len))
      else:
        dataset = torch.utils.data.Subset(dataset, list(range(subset_len)))
    data_loader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=False, **kwargs)
  return data_loader, train_sampler

class AverageMeter(object):
  """Computes and stores the average and current value"""

  def __init__(self, name, fmt=':f'):
    self.name = name
    self.fmt = fmt
    self.reset()

  def reset(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 __str__(self):
    fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
    return fmtstr.format(**self.__dict__)

def accuracy(output, target, topk=(1,)):
  """Computes the accuracy over the k top predictions
    for the specified values of k"""
  with torch.no_grad():
    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

def evaluate(model, val_loader, loss_fn):

  model.eval()
  model = model.cuda()
  top1 = AverageMeter('Acc@1', ':6.2f')
  top5 = AverageMeter('Acc@5', ':6.2f')
  total = 0
  Loss = 0
  for iteraction, (images, labels) in tqdm(
      enumerate(val_loader), total=len(val_loader)):
    images = images.cuda()
    images = images.cuda().to(dtype=torch.get_default_dtype())
    labels = labels.cuda()
    #pdb.set_trace()
    outputs = model(images)
    loss = loss_fn(outputs, labels)
    Loss += loss.item()
    total += images.size(0)
    acc1, acc5 = accuracy(outputs, labels, topk=(1, 5))
    top1.update(acc1[0], images.size(0))
    top5.update(acc5[0], images.size(0))
  return top1.avg, top5.avg, Loss / total

def quantization(title='optimize', model_name='', file_path='', quant_mode=1):

  batch_size = 32

  model = resnet18().cpu()
  model.load_state_dict(torch.load(file_path))

  input = torch.randn([batch_size, 3, 224, 224])
  if quant_mode < 1:
    quant_model = model
  else:
    ## new api
    ####################################################################################
    quantizer = torch_quantizer(
        'calib', model, (input), output_dir='resnet18')

    quant_model = quantizer.quant_model
    #####################################################################################

  # to get loss value after evaluation
  loss_fn = torch.nn.CrossEntropyLoss().cuda()

  val_loader, _ = load_data(
      subset_len=100,
      train=False,
      batch_size=batch_size,
      sample_method='random',
      data_dir=args.data_dir,
      model_name=model_name)

  # record  modules float model accuracy
  # add modules float model accuracy here
  acc_org1 = 0.0
  acc_org5 = 0.0
  loss_org = 0.0

  #register_modification_hooks(model_gen, train=False)
  acc1_gen, acc5_gen, loss_gen = evaluate(quant_model, val_loader, loss_fn)

  # handle quantization result
  if quant_mode > 0:
    quantizer.export_quant_config()
    if quant_mode == 2:
      dump_xmodel('resnet18', True)

  # logging accuracy
  if args.quant_mode == 2:
    basic_info(loss_gen, 'quantized model loss')
    basic_info(acc1_gen, 'quantized model top-1 accuracy')
    basic_info(acc5_gen, 'quantized model top-5 accuracy')
  elif args.quant_mode == 1:
    basic_info(loss_gen, 'calibration model loss')
    basic_info(acc1_gen, 'calibration model top-1 accuracy')
    basic_info(acc5_gen, 'calibration model top-5 accuracy')
  elif args.quant_mode == 0:
    basic_info(loss_gen, 'float model loss')
    basic_info(acc1_gen, 'float model top-1 accuracy')
    basic_info(acc5_gen, 'float model top-5 accuracy')

if __name__ == '__main__':

  model_name = 'resnet18'
  file_path = os.path.join(args.model_dir, model_name + '.pth')

  feature_test = ' float model evaluation'
  if args.quant_mode > 0:
    feature_test = ' quantization'
    # force to merge BN with CONV for better quantization accuracy
    args.optimize = 1
    feature_test += ' with optimization'
  else:
    feature_test = ' float model evaluation'
  title = model_name + feature_test

  print_center_edge(" Start {} test ".format(title))

  # calibration or evaluation
  quantization(
      title=title,
      model_name=model_name,
      file_path=file_path,
      quant_mode=args.quant_mode)

  print_center_edge(" End of {} test ".format(title))
