import numpy as np
import os,sys
import argparse
from tqdm import tqdm
from einops import rearrange, repeat

import torch.nn as nn
import torch
from torch.utils.data import DataLoader
import torch.nn.functional as F
from pytorch_nndct.apis import torch_quantizer, dump_xmodel

from ptflops import get_model_complexity_info


sys.path.append(os.path.join(os.path.dirname(sys.argv[0]), "Uformer"))

import scipy.io as sio
from utils.loader import get_validation_data
import utils

from model import UNet,Uformer,Uformer_Cross,Uformer_CatCross

from skimage import img_as_float32, img_as_ubyte
from skimage.metrics import peak_signal_noise_ratio as psnr_loss
from skimage.metrics import structural_similarity as ssim_loss

parser = argparse.ArgumentParser(description='RGB demoire model test')
parser.add_argument('--input_dir', default='./datasets/demoire/test/',
    type=str, help='Directory of validation images')
parser.add_argument('--result_dir', default='./results/demoire/',
    type=str, help='Directory for results')
parser.add_argument('--weights', default='./models/model_best.pth',
    type=str, help='Path to weights')
parser.add_argument('--quant_mode', type=str, default='calib',
   choices=['calib','test'], help='Quantization mode (calib or test). Default is calib')
parser.add_argument('--quant_model_dir', default='./results/quant_model/',
    type=str, help='Directory for quant model')
parser.add_argument('--gpus', default='0', type=str, help='CUDA_VISIBLE_DEVICES')
parser.add_argument('--arch', default='UNet', type=str, help='arch')
parser.add_argument('--batch_size', default=1, type=int, help='Batch size for dataloader')
parser.add_argument('--save_images', action='store_true', help='Save denoised images in result directory')
parser.add_argument('--embed_dim', type=int, default=32, help='number of data loading workers')    
parser.add_argument('--win_size', type=int, default=8, help='number of data loading workers')
parser.add_argument('--token_projection', type=str,default='linear', help='linear/conv token projection')
parser.add_argument('--token_mlp', type=str,default='leff', help='ffn/leff token mlp')
parser.add_argument('--vit_dim', type=int, default=256, help='vit hidden_dim')
parser.add_argument('--vit_depth', type=int, default=12, help='vit depth')
parser.add_argument('--vit_nheads', type=int, default=8, help='vit hidden_dim')
parser.add_argument('--vit_mlp_dim', type=int, default=512, help='vit mlp_dim')
parser.add_argument('--vit_patch_size', type=int, default=16, help='vit patch_size')
parser.add_argument('--global_skip', action='store_true', default=False, help='global skip connection')
parser.add_argument('--local_skip', action='store_true', default=False, help='local skip connection')
parser.add_argument('--vit_share', action='store_true', default=False, help='share vit module')

parser.add_argument('--train_ps', type=int, default=128, help='patch size of training sample')
args = parser.parse_args()


os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus

utils.mkdir(args.result_dir)

test_dataset = get_validation_data(args.input_dir)
test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False, num_workers=8, drop_last=False)

model_restoration= utils.get_arch(args)
utils.load_checkpoint(model_restoration,args.weights)

model_restoration.cuda()

def evaluate(model, test_loader):
    with torch.no_grad():
        psnr_val_rgb = []
        ssim_val_rgb = []
        for ii, data_test in enumerate(tqdm(test_loader), 0):
            rgb_gt = data_test[0].numpy().squeeze().transpose((1,2,0))
            rgb_noisy = data_test[1].cuda()

            rgb_restored = model(rgb_noisy)
            rgb_restored = torch.clamp(rgb_restored,0,1).cpu().numpy().squeeze().transpose((1,2,0))
            psnr_val_rgb.append(psnr_loss(rgb_restored, rgb_gt))
            ssim_val_rgb.append(ssim_loss(rgb_restored, rgb_gt, multichannel=True))

    psnr_val_rgb = sum(psnr_val_rgb)/len(test_dataset)
    ssim_val_rgb = sum(ssim_val_rgb)/len(test_dataset)
    print("PSNR: %f, SSIM: %f " %(psnr_val_rgb,ssim_val_rgb))
    return psnr_val_rgb,ssim_val_rgb


if (args.quant_mode == 'calib'):

    utils.mkdir(args.result_dir)
    test_dataset = get_validation_data(args.input_dir)
    test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False, num_workers=8, drop_last=False)
    model_restoration= utils.get_arch(args)
    utils.load_checkpoint(model_restoration,args.weights)
    model_restoration.cuda()

    print("===>Finetune with weights: ", args.weights)
    input = torch.randn([args.batch_size, 3, 256, 256])
    quantizer = torch_quantizer('calib', model_restoration, (input),output_dir=args.quant_model_dir) 
    quantized_model = quantizer.quant_model

    print("Evaluate Model before finetune")
    evaluate(quantized_model,test_loader)
    quantizer.export_quant_config()

    # run fast finetuning after quatization 
    print("Run fast finetune with quantized model")
    quantizer.fast_finetune(evaluate,(quantized_model, test_loader))

    print("Evaluate Model after finetune")
    evaluate(quantized_model,test_loader)

    quantizer.export_quant_config()

elif (args.quant_mode == 'test'):

    utils.mkdir(args.result_dir)
    test_dataset = get_validation_data(args.input_dir)
    test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False, num_workers=8, drop_last=False)
    model_restoration= utils.get_arch(args)
    model_restoration.cuda()


    print("===>Test finetuned model with weights: ", args.weights)
    input = torch.randn([args.batch_size, 3, 256, 256])
    quantizer = torch_quantizer('calib', model_restoration, (input), output_dir=args.quant_model_dir)
    quantized_model = quantizer.quant_model
    quantizer.load_ft_param()
    evaluate(quantized_model,test_loader)
    quantizer.export_xmodel(deploy_check=False, output_dir=args.quant_model_dir)