import numpy as np
import torch
import sys
sys.path.append('models')
from SRFlow.code import imread, impad, load_model, t, rgb
from PIL import Image
from torchvision.transforms import PILToTensor

def return_SRFlow_result(lr, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.yml', heat=0.7):
    """
    Apply Super-Resolution using SRFlow model to the input LR (low-resolution) image.

    Args:
    - lr: PIL Image
    - conf_path (str): Configuration file path for the SRFlow model. Default is SRFlow_DF2K_4X.yml.
    - heat (float): Heat parameter for the SRFlow model. Default is 0.6.

    Returns:
    - sr: PIL Image
    """
    model, opt = load_model(conf_path)
    lr = PILToTensor()(lr).permute(1, 2, 0).numpy()

    scale = opt['scale']
    pad_factor = 2

    h, w, c = lr.shape
    lr = impad(lr, bottom=int(np.ceil(h / pad_factor) * pad_factor - h),
                        right=int(np.ceil(w / pad_factor) * pad_factor - w))
    
    lr_t = t(lr)
    heat = opt['heat']
    
    sr_t = model.get_sr(lq=lr_t, heat=heat)

    sr = rgb(torch.clamp(sr_t, 0, 1))
    sr = sr[:h * scale, :w * scale]

    sr = Image.fromarray((sr).astype('uint8'))
    return sr

if __name__ == '__main__':
    ip = Image.open('images/demo.png')
    sr = return_SRFlow_result(ip)
    print(sr.size)