import os

import numpy as np
import torch

from SwinIR.models.network_swinir import SwinIR as net

ROOT_PATH = os.path.dirname(__file__)


class SwinIRDemo:
    def __init__(self):
        self.scale = 4
        self.window_size = 8
        self.tile = 800
        self.tile_overlap = 32
        self.device = 'cuda'
        
        model_path = os.path.join(ROOT_PATH, 'weight/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth')
        self.model = self.model_init(model_path)
    
    def model_init(self, model_path):
        model = net(upscale=self.scale, 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')
        param_key_g = 'params_ema'
        
        pretrained_model = torch.load(model_path)
        model.load_state_dict(
                pretrained_model[param_key_g] if param_key_g in pretrained_model.keys() else pretrained_model,
                strict=True)
        
        model.eval()
        model = model.to(self.device)
        return model
    
    def img_preprocess(self, img_PIL, device, window_size):
        # imgname, img_lq, img_gt = get_image_pair(args, path)  # image to HWC-BGR, float32
        # img_lq = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
        
        # img_lq = img_PIL.convert('BGR')
        img_lq = np.asarray(img_PIL)
        img_lq = img_lq / 255
        
        img_lq = np.transpose(img_lq[:, :, [0, 1, 2]], (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
        
        # 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]
        
        return img_lq, h_old, w_old
    
    def test(self, img_lq):
        b, c, h, w = img_lq.size()
        tile = min(self.tile, h, w)
        assert tile % self.window_size == 0, "tile size should be a multiple of window_size"
        sf = self.scale
        
        stride = tile - self.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)
        W = torch.zeros_like(E)
        
        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 = self.model(in_patch)
                out_patch_mask = torch.ones_like(out_patch)
                
                E[..., h_idx * sf:(h_idx + tile) * sf, w_idx * sf:(w_idx + tile) * sf].add_(out_patch)
                W[..., h_idx * sf:(h_idx + tile) * sf, w_idx * sf:(w_idx + tile) * sf].add_(out_patch_mask)
        output = E.div_(W)
        
        return output
    
    def infer(self, img_lq):
        img_lq, h_old, w_old = self.img_preprocess(img_lq, self.device, self.window_size)
        
        with torch.no_grad():
            output = self.test(img_lq)
            output = output[..., :h_old * self.scale, :w_old * self.scale]
        
        output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
        if output.ndim == 3:
            output = np.transpose(output[[0, 1, 2], :, :], (1, 2, 0))  # CHW-RGB to HCW-BGR
        output = (output * 255.0).round().astype(np.uint8)
        
        return output