import numpy as np
import torch

from code.test import imread, impad, t, load_model, rgb

def return_SRFlow_result(lr_path, conf_path='/models/SRFlow/code/confs/SRFlow_DF2K_4X.yml', heat=0.6):
    model, opt = load_model(conf_path)
    lr = imread(lr_path)

    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]

    return sr