import warnings
warnings.simplefilter('ignore')

import argparse
import cv2
import glob
import os
import shutil
import torch
import sys

sys.path.append(sys.base_prefix + '/codes_bvpp/')

from basicvsrpp_arch import BasicVSRPlusPlus
from data_util import read_img_seq
from img_util import tensor2img

import numpy as np
import torch.nn.functional as F

do_nothing = lambda _: None


def create_kernel():
    gkern2d = np.float32( [
        [1,1,1,1,],
        [1,1,1,1,],
        [1,1,1,1,],
        [1,1,1,1,]] )
    gaussian_kernel = gkern2d / gkern2d.sum()
    zero_kernel = np.zeros_like(gaussian_kernel)
    kernel = np.float32([
        [gaussian_kernel, zero_kernel, zero_kernel],
        [zero_kernel, gaussian_kernel, zero_kernel],
        [zero_kernel, zero_kernel, gaussian_kernel]])
    kernel = torch.from_numpy(kernel)
    return kernel


def downsample_bd(data, kernel, scale, pad_data):
    """
        Note:
            1. `data` should be torch.FloatTensor (data range 0~1) in shape [nchw]
            2. `pad_data` should be enabled in model testing
            3. This function is device agnostic, i.e., data/kernel could be on cpu or gpu
    """
    if pad_data:
        data = F.pad(data, (1, 3, 1, 3), 'reflect')
    data = F.conv2d(data, kernel, stride=scale, bias=None, padding=0)
    return data


def inference2(auto_cast, imgs, imgnames, model, save_path, extension, device, scale):
    #with torch.no_grad():
    with torch.inference_mode():
        with torch.autocast(device.type, enabled=auto_cast):
            outputs = model(imgs)

    # save imgs
    outputs = outputs.squeeze()
    outputs = list(outputs)
    for output, imgname in zip(outputs, imgnames):
        blur_kernel = create_kernel().to(device)
        output = downsample_bd(output, blur_kernel, scale, pad_data=True)
        output = tensor2img(output)
        cv2.imwrite(os.path.join(save_path, f'{imgname}.{extension}'), output)
        print('\rwrite:' + imgname, end='')


def inference(auto_cast, imgs, imgnames, model, save_path, extension, device, reduce_type, reduce_scale):
    #with torch.no_grad():
    with torch.inference_mode():
        with torch.autocast(device.type, enabled=auto_cast):
            outputs = model(imgs)

    # save imgs
    outputs = outputs.squeeze()
    outputs = list(outputs)
    for output, imgname in zip(outputs, imgnames):
        output = tensor2img(output)
        if reduce_type > 0:
            output = cv2.resize(output, None, fx=reduce_scale, fy=reduce_scale, interpolation=reduce_type)
        cv2.imwrite(os.path.join(save_path, f'{imgname}.{extension}'), output)
        print('\rwrite:' + imgname, end='')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("model_path")
    parser.add_argument("input_path")
    parser.add_argument("save_path")
    parser.add_argument("interval", type=int)
    parser.add_argument("device", type=int)
    parser.add_argument("fraction", type=int)
    #parser.add_argument('--channel')
    #parser.add_argument('--block')
    parser.add_argument("extension")
    parser.add_argument("-ac", action="store_true")
    parser.add_argument("-r", "--reduce_type", type=int, default=0)
    args = parser.parse_args()

    # select device
    if torch.cuda.is_available() and args.device < 7:
        device = torch.device('cuda', args.device)
        print('GPU_ID=' + str(args.device))
        #print('device=' + torch.cuda.get_device_name(device))
        if args.fraction < 100:
            torch.cuda.set_per_process_memory_fraction(args.fraction*0.01, args.device)
        print('VRAM_fraction=' + str(args.fraction) + '%')
        device_name = torch.cuda.get_device_name(device)
        print('device_name =', device_name)
        if 'ZLUDA' in device_name:
            torch.backends.cudnn.enabled = False
            torch.backends.cuda.enable_flash_sdp(False)
            torch.backends.cuda.enable_flash_sdp = do_nothing
            torch.backends.cuda.enable_math_sdp(True)
            torch.backends.cuda.enable_math_sdp = do_nothing
            torch.backends.cuda.enable_mem_efficient_sdp(False)
            torch.backends.cuda.enable_mem_efficient_sdp = do_nothing
            if hasattr(torch.backends.cuda, "enable_cudnn_sdp"):
                torch.backends.cuda.enable_cudnn_sdp(False)
                torch.backends.cuda.enable_cudnn_sdp = do_nothing
        #else:
        #    torch.backends.cudnn.enabled = False
        #    if hasattr(torch.backends.cuda, "enable_cudnn_sdp"):
        #        torch.backends.cuda.enable_cudnn_sdp(False)
        #        torch.backends.cuda.enable_cudnn_sdp = do_nothing
    else:
        #device = torch.device('cpu')
        #print('GPU not used')
        print('This program requires CUDA.')
        sys.exit()

    scale = 8
    if '8X' in args.model_path:
        scale = 8
    if '4X' in args.model_path:
        scale = 4

    # set up model
    c = 64
    b = 7
    if(args.model_path[-7] == '@'):
        c = 1 << int(args.model_path[-6])
        b = int(args.model_path[-5], base=16)
    print('channel=' + str(c))
    print('block=' + str(b)) 
    print('model=' + args.model_path);
    print('input=' + args.input_path);
    print('extension=' + args.extension)
    if args.reduce_type > 0:
        print('reduce=' + str(args.reduce_type))
    model = BasicVSRPlusPlus(mid_channels=c, num_blocks=b, scale=scale)
    model.load_state_dict(torch.load(sys.base_prefix + '\\' + args.model_path, map_location=device)['params'], strict=False)

    """
    checkpoint = torch.load(args.model_path, map_location=device)['params']
    for key in list(checkpoint.keys()):
        print('key=', key)
        if 'deform_align.' in key:
            checkpoint[key.replace('deform_align.', '')] = checkpoint[key]
            del checkpoint[key]
    model.load_state_dict(checkpoint)
    """

    model.eval()
    model = model.to(device)

    auto_cast = args.ac and device != torch.device('cpu')
    print('auto_cast=' + str(auto_cast))

    os.makedirs(args.save_path, exist_ok=True)

    # load data and inference
    imgs_list = sorted(glob.glob(os.path.join(args.input_path, '*.png')))
    num_imgs = len(imgs_list)
    if num_imgs <= args.interval:  # too many images may cause CUDA out of memory
        imgs, imgnames = read_img_seq(imgs_list, require_mod_crop=False, scale=1, return_imgname=True)
        height, width, channels = imgs[0].shape
        imgs = imgs.unsqueeze(0).to(device)
        print('inference:0-' + str(num_imgs-1))
        if args.reduce_type == 5:
            inference2(auto_cast, imgs, imgnames, model, args.save_path, args.extension, device, scale)
        else:
            inference(auto_cast, imgs, imgnames, model, args.save_path, args.extension, device, args.reduce_type, 1.0 / scale)
        print('')
    else:
        for idx in range(0, num_imgs, args.interval):
            interval = min(args.interval, num_imgs - idx)
            imgs, imgnames = read_img_seq(imgs_list[idx:idx + interval], require_mod_crop=False, scale=1, return_imgname=True)
            height, width, channels = imgs[0].shape
            imgs = imgs.unsqueeze(0).to(device)
            print('inference:' + str(idx) + '-' + str(idx+interval-1))
            if args.reduce_type == 5:
                inference2(auto_cast, imgs, imgnames, model, args.save_path, args.extension, device, scale)
            else:
                inference(auto_cast, imgs, imgnames, model, args.save_path, args.extension, device, args.reduce_type, 1.0 / scale)
            print('')


if __name__ == '__main__':
    main()