from functools import cache

import torch
import torch.nn.functional as F
from einops import reduce
from jaxtyping import Float
from lpips import LPIPS
from torch import Tensor


@torch.no_grad()
def compute_psnr(
    ground_truth: Float[Tensor, "batch channel height width"],
    predicted: Float[Tensor, "batch channel height width"],
) -> Float[Tensor, " batch"]:
    with torch.autocast(device_type=predicted.device.type, enabled=False):
        ground_truth = ground_truth.float().clip(min=0, max=1)
        predicted = predicted.float().clip(min=0, max=1)
        mse = reduce((ground_truth - predicted) ** 2, "b c h w -> b", "mean")
        return -10 * mse.log10()


@cache
def get_lpips(device: torch.device) -> LPIPS:
    model = LPIPS(net="vgg").eval().requires_grad_(False)
    return model.to(device=device, dtype=torch.float32)


@torch.no_grad()
def compute_lpips(
    ground_truth: Float[Tensor, "batch channel height width"],
    predicted: Float[Tensor, "batch channel height width"],
) -> Float[Tensor, " batch"]:
    with torch.autocast(device_type=predicted.device.type, enabled=False):
        value = get_lpips(predicted.device).forward(
            ground_truth.float(),
            predicted.float(),
            normalize=True,
        )
        return value[:, 0, 0, 0].float()


@torch.no_grad()
def compute_ssim(
    ground_truth: Float[Tensor, "batch channel height width"],
    predicted: Float[Tensor, "batch channel height width"],
    *,
    win_size: int = 11,
    gaussian_weights: bool = True,
    data_range: float = 1.0,
    use_sample_covariance: bool = True,
    sigma: float = 1.5,
    k1: float = 0.01,
    k2: float = 0.03,
) -> Float[Tensor, " batch"]:
    if ground_truth.shape != predicted.shape or ground_truth.ndim != 4:
        raise ValueError("SSIM inputs must have the same BCHW shape.")
    if win_size <= 0 or win_size % 2 == 0:
        raise ValueError("win_size must be a positive odd integer.")
    if min(ground_truth.shape[-2:]) < win_size:
        raise ValueError("win_size must not exceed the image dimensions.")
    if gaussian_weights and sigma <= 0:
        raise ValueError("sigma must be positive when using Gaussian weights.")
    if data_range <= 0 or k1 < 0 or k2 < 0:
        raise ValueError("data_range must be positive and k1/k2 non-negative.")

    with torch.autocast(device_type=predicted.device.type, enabled=False):
        ground_truth = ground_truth.float()
        predicted = predicted.float()

        if gaussian_weights:
            coordinates = torch.arange(
                win_size,
                dtype=torch.float32,
                device=predicted.device,
            )
            coordinates -= (win_size - 1) / 2
            kernel_1d = torch.exp(-(coordinates.square()) / (2 * sigma**2))
            kernel_1d /= kernel_1d.sum()
            kernel = torch.outer(kernel_1d, kernel_1d)
        else:
            kernel = torch.full(
                (win_size, win_size),
                1 / win_size**2,
                dtype=torch.float32,
                device=predicted.device,
            )

        channels = predicted.shape[1]
        moments = torch.cat(
            (
                ground_truth,
                predicted,
                ground_truth.square(),
                predicted.square(),
                ground_truth * predicted,
            ),
            dim=1,
        )
        kernels = kernel[None, None].expand(5 * channels, 1, -1, -1)
        ux, uy, uxx, uyy, uxy = F.conv2d(
            moments,
            kernels,
            groups=5 * channels,
        ).chunk(5, dim=1)

        covariance_normalization = 1.0
        if use_sample_covariance:
            num_pixels = win_size**2
            if num_pixels == 1:
                raise ValueError(
                    "Sample covariance requires a window larger than one pixel."
                )
            covariance_normalization = num_pixels / (num_pixels - 1)

        vx = covariance_normalization * (uxx - ux.square())
        vy = covariance_normalization * (uyy - uy.square())
        vxy = covariance_normalization * (uxy - ux * uy)

        c1 = (k1 * data_range) ** 2
        c2 = (k2 * data_range) ** 2
        ssim = (
            (2 * ux * uy + c1)
            * (2 * vxy + c2)
            / ((ux.square() + uy.square() + c1) * (vx + vy + c2))
        )
        ssim = ssim.mean(dim=(1, 2, 3))

    return ssim