import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from typing import Union
import itertools
import numpy as np


def load_satlas_sr(device: Union[str, torch.device] = "cuda") -> RRDBNet:
    # Load the weights
    weights_file = "satlas/weights/esrgan_1S2.pth"
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Create the model
    model = RRDBNet(
        num_in_ch=3,
        num_out_ch=3,
        num_feat=64,
        num_block=23,
        num_grow_ch=32,
        scale=4
    ).to(device)

    # Setup the weights
    state_dict = torch.load(weights_file)
    model.load_state_dict(state_dict['params_ema'])
    model.eval()

    # no gradients
    for param in model.parameters():
        param.requires_grad = False

    return model


def run_satlas(
    model: RRDBNet,
    lr: torch.Tensor,
    hr: torch.Tensor,
    cropsize: int = 32,
    overlap: int = 0,
    device: Union[str, torch.device] = "cuda"
) -> torch.Tensor:
    # Load the LR image 
    lr = torch.from_numpy(lr[[3, 2, 1]]/3558).float().to(device).clamp(0, 1)

    # Select the raster with the lowest resolution
    tshp = lr.shape
    
    # if the image is too small, return (0, 0)
    if (tshp[1] < cropsize) and (tshp[2] < cropsize):
        return [(0, 0)]

    # Define relative coordinates.
    xmn, xmx, ymn, ymx = (0, tshp[1], 0, tshp[2])

    if overlap > cropsize:
        raise ValueError("The overlap must be smaller than the cropsize")

    xrange = np.arange(xmn, xmx, (cropsize - overlap))
    yrange = np.arange(ymn, ymx, (cropsize - overlap))

    # If there is negative values in the range, change them by zero.
    xrange[xrange < 0] = 0
    yrange[yrange < 0] = 0

    # Remove the last element if it is outside the tensor
    xrange = xrange[xrange - (tshp[1] - cropsize) <= 0]
    yrange = yrange[yrange - (tshp[2] - cropsize) <= 0]

    # If the last element is not (tshp[1] - cropsize) add it!
    if xrange[-1] != (tshp[1] - cropsize):
        xrange = np.append(xrange, tshp[1] - cropsize)
    if yrange[-1] != (tshp[2] - cropsize):
        yrange = np.append(yrange, tshp[2] - cropsize)

    # Create all the relative coordinates
    mrs = list(itertools.product(xrange, yrange))

    # Predict the image
    sr = torch.zeros(3, tshp[1]*4, tshp[2]*4)
    for x, y in mrs:
        crop = lr[:, x:x+cropsize, y:y+cropsize]        
        sr_crop = model(crop[None])[0]
        sr[:, x*4:(x+cropsize)*4, y*4:(y+cropsize)*4] = sr_crop
    
    # Save the result
    results = {
        "lr": (lr.cpu().numpy() * 3558).astype(np.uint16),
        "sr": (sr.cpu().numpy() * 3558).astype(np.uint16),
        "hr": hr[0:3]
    }
    

    return results