import gradio as gr
import os
import numpy as np
import torch
from models.network_swinir import SwinIR


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device: %s" % device)
default_models = {
    "sr": "weights/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth",
    "denoise": "weights/005_colorDN_DFWB_s128w8_SwinIR-M_noise25.pth"
    }
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True


denoise_model = SwinIR(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').to(device)
param_key_g = 'params'
try:
    pretrained_model = torch.load(default_models["denoise"])
    denoise_model.load_state_dict(pretrained_model[param_key_g] if param_key_g in pretrained_model.keys() else pretrained_model, strict=True)
except: print("Loading model failed")
denoise_model.eval()

sr_model = SwinIR(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').to(device)
param_key_g = 'params_ema'
try:
    pretrained_model = torch.load(default_models["sr"])
    sr_model.load_state_dict(pretrained_model[param_key_g] if param_key_g in pretrained_model.keys() else pretrained_model, strict=True)
except: print("Loading model failed")
sr_model.eval()


def sr(input_img):

    window_size = 8
        # read image
    img_lq = input_img.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)  # CHW-RGB to NCHW-RGB

    # 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 = sr_model(img_lq)
        output = output[..., :h_old * 4, :w_old * 4]

    # save image
    output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
    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

    return output

def denoise(input_img):

    window_size = 8
    # read image
    img_lq = input_img.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)  # CHW-RGB to NCHW-RGB

    # 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 = denoise_model(img_lq)
        output = output[..., :h_old * 4, :w_old * 4]

    # save image
    output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
    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
    
    return output

title = " AISeed AI Application Demo "
description = "# A Demo of Deep Learning for Image Restoration"
example_list = [["examples/" + example] for example in os.listdir("examples")]

with gr.Blocks() as demo:
    demo.title = title
    gr.Markdown(description)
    with gr.Row():
        with gr.Column():
            im = gr.Image(label="Input Image")
            im_2 = gr.Image(label="Enhanced Image")
            
        with gr.Column():
           
            btn1 = gr.Button(value="Enhance Resolution")
            btn1.click(sr, inputs=[im], outputs=[im_2])
            btn2 = gr.Button(value="Denoise")
            btn2.click(denoise, inputs=[im], outputs=[im_2])
            gr.Examples(examples=example_list,
                inputs=[im],
                outputs=[im_2])
    
if __name__ == "__main__":
    demo.launch()