import torch
import torch.nn.functional as F
from skimage.metrics import peak_signal_noise_ratio


def pad_to_factor(x, factor=8):
    if factor <= 1:
        return x, (0, 0)

    h, w = x.shape[-2:]
    pad_h = (factor - h % factor) % factor
    pad_w = (factor - w % factor) % factor
    if pad_h == 0 and pad_w == 0:
        return x, (h, w)

    return F.pad(x, (0, pad_w, 0, pad_h), mode="reflect"), (h, w)


def crop_to_size(x, size):
    h, w = size
    return x[..., :h, :w]


def _tta_inputs(x):
    return [
        (x, lambda y: y),
        (torch.flip(x, dims=[-1]), lambda y: torch.flip(y, dims=[-1])),
        (torch.flip(x, dims=[-2]), lambda y: torch.flip(y, dims=[-2])),
        (torch.rot90(x, k=1, dims=(-2, -1)), lambda y: torch.rot90(y, k=3, dims=(-2, -1))),
    ]


def _unwrap_output(output):
    if isinstance(output, tuple):
        return output[0]
    return output


@torch.no_grad()
def run_model(model, x, pad_factor=8, tta=False):
    if not tta:
        padded, size = pad_to_factor(x, pad_factor)
        out = _unwrap_output(model(padded))
        return crop_to_size(out, size)

    outputs = []
    for aug_x, deaug in _tta_inputs(x):
        padded, size = pad_to_factor(aug_x, pad_factor)
        out = crop_to_size(_unwrap_output(model(padded)), size)
        outputs.append(deaug(out))
    return torch.stack(outputs, dim=0).mean(dim=0)


def batch_rgb_psnr(pred, target):
    pred_np = pred.detach().cpu().numpy()
    target_np = target.detach().cpu().numpy()
    total = 0.0
    for pred_img, target_img in zip(pred_np, target_np):
        total += peak_signal_noise_ratio(target_img, pred_img, data_range=1.0)
    return total / max(pred_np.shape[0], 1)