import argparse
import cv2
import glob
import numpy as np
from collections import OrderedDict
import os
import torch
import requests
import math


from utils import utils_image as util
import pdb

from models.model_gan import DeepTen, get_deepten_resnet50_minc
import torch.nn as nn

import lpips

from basicsr.metrics import calculate_psnr, calculate_ssim
from basicsr.utils import bgr2ycbcr, scandir

"""
def calculate_psnr(src,gt):
    # src : return value of cv2.imread(src_path)
    # gt : return value of cv2.imread(gt_path)
	
	height, width, channel = src.shape
	size = height*width

	sb,sg,sr = cv2.split(src)
	rb,rg,rr = cv2.split(gt)

	mseb = ((sb-rb)**2).sum()
	mseg = ((sg-rg)**2).sum()
	mser = ((sr-rr)**2).sum()

	MSE = (mseb+mseg+mser)/(3*size)
	psnr = 10*math.log10(255**2/MSE)
	return psnr

def calculate_ssim(img1, img2):
    # img1 : return value of cv2.imread(img1_path)
    # img2 : return value of cv2.imread(img2_path)
	
	C1 = (0.01 * 255)**2
	C2 = (0.03 * 255)**2
	img1 = img1.astype(np.float64)
	img2 = img2.astype(np.float64)
	kernel = cv2.getGaussianKernel(11, 1.5)
	window = np.outer(kernel, kernel.transpose())
	mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
	mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
	mu1_sq = mu1**2
	mu2_sq = mu2**2
	mu1_mu2 = mu1 * mu2
	sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
	sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
	sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
	ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *(sigma1_sq + sigma2_sq + C2))
	return ssim_map.mean()
"""

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--network', type=str, default='SwinIR', help='SwinIR | BSR')
    parser.add_argument('--task', type=str, default='color_dn', help='classical_sr, lightweight_sr, real_sr, '
                                                                     'gray_dn, color_dn, jpeg_car')
    parser.add_argument('--scale', type=int, default=1, help='scale factor: 1, 2, 3, 4, 8') # 1 for dn and jpeg car
    parser.add_argument('--noise', type=int, default=15, help='noise level: 15, 25, 50')
    parser.add_argument('--jpeg', type=int, default=40, help='scale factor: 10, 20, 30, 40')
    parser.add_argument('--training_patch_size', type=int, default=128, help='patch size used in training SwinIR. '
                                       'Just used to differentiate two different settings in Table 2 of the paper. '
                                       'Images are NOT tested patch by patch.')
    parser.add_argument('--large_model', action='store_true', help='use large model, only provided for real image sr')
    parser.add_argument('--model_path', type=str,
                        default='model_zoo/swinir/001_classicalSR_DIV2K_s48w8_SwinIR-M_x2.pth')
    parser.add_argument('--model_DT_path', type=str,
                        default=None)
    parser.add_argument('--print_img', type=bool,
                        default=False)
    parser.add_argument('--folder_lq', type=str, default=None, help='input low-quality test image folder')
    parser.add_argument('--folder_gt', type=str, default=None, help='input ground-truth test image folder')
    parser.add_argument('--tile', type=int, default=None, help='Tile size, None for no tile during testing (testing as a whole)')
    parser.add_argument('--tile_overlap', type=int, default=32, help='Overlapping of different tiles')
    args = parser.parse_args()
    
    
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # set up model
    if os.path.exists(args.model_path):
        print(f'loading model from {args.model_path}')
    else:
        os.makedirs(os.path.dirname(args.model_path), exist_ok=True)
        url = 'https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/{}'.format(os.path.basename(args.model_path))
        r = requests.get(url, allow_redirects=True)
        print(f'downloading model {args.model_path}')
        open(args.model_path, 'wb').write(r.content)
    
        
        
    model,model_DT = define_model(args)
    model.eval()
    model = model.to(device)
    
    if model_DT is not None:
        model_DT.eval()
        model_DT = model_DT.to(device)

    # setup folder and path
    folder, save_dir, border, window_size = setup(args)
    os.makedirs(save_dir, exist_ok=True)
    
    
    folder_lq = '/mnt/server5_hard1/seungjun/KAIR-master/testsets/DIV2K4D/LR_x4/'
    parts=['part0','part1', 'part2', 'part3', 'part4']
    #parts=['part2', 'part3', 'part4']
    folder_gt = '/mnt/server5_hard1/seungjun/KAIR-master/testsets/DIV2K4D/HR/'
    
    loss_fn_alex = lpips.LPIPS(net='alex') # best forward scores
    loss_fn_vgg = lpips.LPIPS(net='vgg') # closer to "traditional" perceptual loss, when used for optimization

    
    for part in parts:
        
        if part == 'part0':
            tmp_folder_lq=folder_lq
            tmp_folder_gt=folder_gt
            folder_lr='/mnt/server5_hard1/seungjun/KAIR-master/testsets/testL/X4/'
            folder_gt='/mnt/server5_hard1/seungjun/KAIR-master/testsets/testH/'
        else:
            folder_lq=tmp_folder_lq
            folder_gt=tmp_folder_gt
            folder_lr=folder_lq+part+'/'
        test_results = OrderedDict()
        test_results['psnr'] = []
        test_results['ssim'] = []
        test_results['psnr_y'] = []
        test_results['ssim_y'] = []
        test_results['psnr_b'] = []
        test_results['lpips_d'] = []
        psnr, ssim, psnr_y, ssim_y, psnr_b, lpips_d = 0, 0, 0, 0, 0, 0
        for idx, path in enumerate(sorted(glob.glob(os.path.join(folder_lr, '*')))):
            # read image
            imgname, img_lq, _ = get_image_pair(args, path)  # image to HWC-BGR, float32
            if part in ['part2','part3']:
                img_number=imgname[3:]
                img_number=800+int(img_number)
                imgname='0'+str(img_number)
            img_path_gt=folder_gt+imgname+'.png'
            img_gt = cv2.imread(img_path_gt, cv2.IMREAD_COLOR)
            img_gt_npy=img_gt.astype(np.float32) / 255.
            img_lq = np.transpose(img_lq if img_lq.shape[2] == 1 else img_lq[:, :, [2, 1, 0]], (2, 0, 1))  # HCW-BGR to CHW-RGB
            img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(device)
            
            img_gt_npy = np.transpose(img_gt_npy if img_gt_npy.shape[2] == 1 else img_gt_npy[:, :, [2, 1, 0]], (2, 0, 1))  # HCW-BGR to CHW-RGB
            img_gt_npy = torch.from_numpy(img_gt_npy).float().unsqueeze(0)  

            # inference
            with torch.no_grad():
                # pad input image to be a multiple of window_size
                _, _, h_old, w_old = img_lq.size()
                h_pad = (h_old // window_size + 1) * window_size - h_old
                w_pad = (w_old // window_size + 1) * window_size - w_old
                img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
                img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
                output, att,E_out = test(img_lq, model, model_DT, args, window_size)
                output = output[..., :h_old * args.scale, :w_old * args.scale]
                
                if model_DT is not None:
                    att = att[..., :h_old * args.scale, :w_old * args.scale]
                    E_out = E_out[..., :h_old * args.scale, :w_old * args.scale]
                    E_out=E_out.data.squeeze().float().cpu().clamp_(0, 1).numpy()

            # save image
            output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
            
            if not os.path.exists(f'{save_dir}/{part}'):
      
                # if the demo_folder directory is not present 
                # then create it.
                os.makedirs(f'{save_dir}/{part}')
                
            if img_gt is not None:
                
                lpips_d=loss_fn_alex(torch.Tensor(img_gt_npy)*2.0-1.0, torch.Tensor(output).unsqueeze(0)*2.0-1.0)
                lpips_d=lpips_d.item()
                test_results['lpips_d'].append(lpips_d)    
                
            if output.ndim == 3:
                    output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))  # CHW-RGB to HCW-BGR
                    output = (output * 255.0).round().astype(np.uint8)  # float32 to uint8
                    
            if model_DT is not None:
                att = att.data.squeeze().float().cpu().clamp_(0, 1).numpy()
                E_out = np.transpose(E_out[[2, 1, 0], :, :], (1, 2, 0))
                E_out = (E_out * 255.0).round().astype(np.uint8)  # float32 to uint8
                if args.print_img:
                    cv2.imwrite(f'{save_dir}/{part}/{imgname}_att.png', att)
                    cv2.imwrite(f'{save_dir}/{part}/{imgname}_Eout.png', E_out)
                    cv2.imwrite(f'{save_dir}/{part}/{imgname}_SAITGAN.png', output)
            else:
                if args.print_img:
                    cv2.imwrite(f'{save_dir}/{part}/{imgname}_SWINIR.png', output)
            
            # evaluate psnr/ssim/psnr_b
            if img_gt is not None:                
                                
                #psnr = util.calculate_psnr(output, img_gt, border=border)
                ssim = calculate_ssim(output, img_gt, input_order='HWC', crop_border=0)
                psnr= calculate_psnr(output, img_gt,input_order='HWC', crop_border=0)
                
                test_results['psnr'].append(psnr)
                test_results['ssim'].append(ssim)
                
                if args.task in ['jpeg_car']:
                    psnr_b = util.calculate_psnrb(output, img_gt, border=border)
                    test_results['psnr_b'].append(psnr_b)
                print('Testing {:d} {:20s} - PSNR: {:.2f} dB; SSIM: {:.4f}; '
                      'PSNR_B: {:.2f} dB; LPIPS_D: {:.4f}:.'.
                      format(idx, imgname, psnr, ssim, psnr_b,lpips_d), flush=True)
            else:
                print('Testing {:d} {:20s}'.format(idx, imgname))

    # summarize psnr/ssim
        if img_gt is not None:
            ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
            ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])
            ave_lpips_d = sum(test_results['lpips_d']) / len(test_results['lpips_d'])
            print('\n{} \n-- Average PSNR/SSIM(RGB)/LPIPS-D: {:.2f} dB; {:.4f}; {:.4f}'.format(save_dir, ave_psnr, ave_ssim, ave_lpips_d))
            
            if args.task in ['jpeg_car']:
                ave_psnr_b = sum(test_results['psnr_b']) / len(test_results['psnr_b'])
                print('-- Average PSNR_B: {:.2f} dB'.format(ave_psnr_b))


def define_model(args):
    # 001 classical image sr
    if args.network=='SwinIR':
        from models.network_swinir import SwinIR as net
        if args.task == 'classical_sr':
            model = net(upscale=args.scale, in_chans=3, img_size=args.training_patch_size, window_size=8,
                        img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
                        mlp_ratio=2, upsampler='pixelshuffle', resi_connection='1conv')
            param_key_g = 'params'

        # 002 lightweight image sr
        # use 'pixelshuffledirect' to save parameters
        elif args.task == 'lightweight_sr':
            model = net(upscale=args.scale, in_chans=3, img_size=64, window_size=8,
                        img_range=1., depths=[6, 6, 6, 6], embed_dim=60, num_heads=[6, 6, 6, 6],
                        mlp_ratio=2, upsampler='pixelshuffledirect', resi_connection='1conv')
            param_key_g = 'params'

        # 003 real-world image sr
        elif args.task == 'real_sr':
            if not args.large_model:
                # use 'nearest+conv' to avoid block artifacts
                model = net(upscale=4, in_chans=3, img_size=64, window_size=8,
                            img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
                            mlp_ratio=2, upsampler='nearest+conv', resi_connection='1conv')
            else:
                # larger model size; use '3conv' to save parameters and memory; use ema for GAN training
                model = net(upscale=4, in_chans=3, img_size=64, window_size=8,
                            img_range=1., depths=[6, 6, 6, 6, 6, 6, 6, 6, 6], embed_dim=240,
                            num_heads=[8, 8, 8, 8, 8, 8, 8, 8, 8],
                            mlp_ratio=2, upsampler='nearest+conv', resi_connection='3conv')
            param_key_g = 'params_ema'

        # 004 grayscale image denoising
        elif args.task == 'gray_dn':
            model = net(upscale=1, in_chans=1, img_size=128, window_size=8,
                        img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
                        mlp_ratio=2, upsampler='', resi_connection='1conv')
            param_key_g = 'params'

        # 005 color image denoising
        elif args.task == 'color_dn':
            model = net(upscale=1, in_chans=3, img_size=128, window_size=8,
                        img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
                        mlp_ratio=2, upsampler='', resi_connection='1conv')
            param_key_g = 'params'

        # 006 JPEG compression artifact reduction
        # use window_size=7 because JPEG encoding uses 8x8; use img_range=255 because it's sligtly better than 1
        elif args.task == 'jpeg_car':
            model = net(upscale=1, in_chans=1, img_size=126, window_size=7,
                        img_range=255., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
                        mlp_ratio=2, upsampler='', resi_connection='1conv')
            param_key_g = 'params'
            
        pretrained_model = torch.load(args.model_path)    
        model.load_state_dict(pretrained_model[param_key_g] if param_key_g in pretrained_model.keys() else pretrained_model, strict=True)
        
    elif args.network=='BSR':
        from models.network_rrdbnet import RRDBNet as net
        model = net(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=4)  # define network  , sf= scale factor of SR
        pretrained_model = torch.load(args.model_path)    
        model.load_state_dict(pretrained_model, strict=True)
        
    elif args.network=='SRFormer':
        from basicsr.archs.srformer_arch import SRFormer as net
        model = net(upscale=4, in_chans=3, img_size=64, window_size=22,
                            img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180,
                            num_heads=[6,6,6,6,6,6],
                            mlp_ratio=2, upsampler='pixelshuffle', resi_connection='1conv')
        pretrained_model = torch.load(args.model_path)    
        model.load_state_dict(pretrained_model, strict=True)
        
    if args.model_DT_path is not None:
        model_kwargs = {'pretrained': True}
        model_DT = get_deepten_resnet50_minc(**model_kwargs) 
        pretrained_DT = torch.load(args.model_DT_path)
        model_DT.load_state_dict(pretrained_DT)
    else:
        model_DT = None
    
    
    
        
    return model, model_DT


def setup(args):
    # 001 classical image sr/ 002 lightweight image sr
    if args.task in ['classical_sr', 'lightweight_sr']:
        save_dir = f'results/swinir_{args.task}_x{args.scale}'
        folder = args.folder_gt
        border = args.scale
        window_size = 8

    # 003 real-world image sr
    elif args.task in ['real_sr']:
        if args.model_DT_path is not None:
            save_dir = f'results/Ours_{args.network}_DeepTen3_DIV2K4D_{args.task}_x{args.scale}'
        else:
            save_dir = f'results/Naive_{args.network}_DIV2K4D_{args.task}_x{args.scale}'
        if args.large_model:
            save_dir += '_large'
        folder = args.folder_lq
        border = 0
        window_size = 8

    # 004 grayscale image denoising/ 005 color image denoising
    elif args.task in ['gray_dn', 'color_dn']:
        save_dir = f'results/swinir_{args.task}_noise{args.noise}'
        folder = args.folder_gt
        border = 0
        window_size = 8

    # 006 JPEG compression artifact reduction
    elif args.task in ['jpeg_car']:
        save_dir = f'results/swinir_{args.task}_jpeg{args.jpeg}'
        folder = args.folder_gt
        border = 0
        window_size = 7

    return folder, save_dir, border, window_size


def get_image_pair(args, path):
    (imgname, imgext) = os.path.splitext(os.path.basename(path))

    # 001 classical image sr/ 002 lightweight image sr (load lq-gt image pairs)
    if args.task in ['classical_sr', 'lightweight_sr']:
        img_gt = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
        img_lq = cv2.imread(f'{args.folder_lq}/{imgname}x{args.scale}{imgext}', cv2.IMREAD_COLOR).astype(
            np.float32) / 255.

    # 003 real-world image sr (load lq image only)
    elif args.task in ['real_sr']:
        img_gt = None
        img_lq = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.

    # 004 grayscale image denoising (load gt image and generate lq image on-the-fly)
    elif args.task in ['gray_dn']:
        img_gt = cv2.imread(path, cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255.
        np.random.seed(seed=0)
        img_lq = img_gt + np.random.normal(0, args.noise / 255., img_gt.shape)
        img_gt = np.expand_dims(img_gt, axis=2)
        img_lq = np.expand_dims(img_lq, axis=2)

    # 005 color image denoising (load gt image and generate lq image on-the-fly)
    elif args.task in ['color_dn']:
        img_gt = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
        np.random.seed(seed=0)
        img_lq = img_gt + np.random.normal(0, args.noise / 255., img_gt.shape)

    # 006 JPEG compression artifact reduction (load gt image and generate lq image on-the-fly)
    elif args.task in ['jpeg_car']:
        img_gt = cv2.imread(path, 0)
        result, encimg = cv2.imencode('.jpg', img_gt, [int(cv2.IMWRITE_JPEG_QUALITY), args.jpeg])
        img_lq = cv2.imdecode(encimg, 0)
        img_gt = np.expand_dims(img_gt, axis=2).astype(np.float32) / 255.
        img_lq = np.expand_dims(img_lq, axis=2).astype(np.float32) / 255.

    return imgname, img_lq, img_gt


def test(img_lq, model, model_DT, args, window_size):
    m_up=nn.Upsample(scale_factor=4, mode='bilinear')
    if args.tile is None:
        # test the image as a whole
        E_out = model(img_lq)
        if model_DT is not None:
            att = model_DT(img_lq)
            L_up=m_up(img_lq)
            att=m_up(att)
            output = att * E_out + (1 - att) * L_up
        else:
            output=E_out
            att=None
    else:
        # test the image tile by tile
        b, c, h, w = img_lq.size()
        
        tile = min(args.tile, h, w)
        assert tile % window_size == 0, "tile size should be a multiple of window_size"
        tile_overlap = args.tile_overlap
        sf = args.scale

        stride = tile - tile_overlap
        h_idx_list = list(range(0, h-tile, stride)) + [h-tile]
        w_idx_list = list(range(0, w-tile, stride)) + [w-tile]
        E = torch.zeros(b, c, h*sf, w*sf).type_as(img_lq)
        Att_ = torch.zeros(b, 1, h*sf, w*sf).type_as(img_lq)
        W = torch.zeros_like(E)
        W_att = torch.zeros_like(Att_)

        for h_idx in h_idx_list:
            for w_idx in w_idx_list:
                in_patch = img_lq[..., h_idx:h_idx+tile, w_idx:w_idx+tile]
                out_patch = model(in_patch)
                #print('out patch size is {}'.format(out_patch.shape))
                att_patch = model_DT(in_patch)
                L_up=m_up(in_patch)
                out_patch = att_patch * out_patch + (1 - att_patch) * L_up
                out_patch_mask = torch.ones_like(out_patch)
                att_patch_mask = torch.ones_like(att_patch)

                E[..., h_idx*sf:(h_idx+tile)*sf, w_idx*sf:(w_idx+tile)*sf].add_(out_patch)
                Att_[..., h_idx*sf:(h_idx+tile)*sf, w_idx*sf:(w_idx+tile)*sf].add_(att_patch)
                W[..., h_idx*sf:(h_idx+tile)*sf, w_idx*sf:(w_idx+tile)*sf].add_(out_patch_mask)
                W_att[..., h_idx*sf:(h_idx+tile)*sf, w_idx*sf:(w_idx+tile)*sf].add_(att_patch_mask)
                
        output = E.div_(W)
        att = Att_.div_(W_att)
        E_out=None

    return output, att, E_out

if __name__ == '__main__':
    main()