"""PSNR / SSIM / LPIPS between a strategy's frames and its matched FFFF frames.

Computed on the decoded RGB tensors in [0, 1], before MP4 encoding, as the
protocol requires (section 8.1) -- video compression noise would otherwise leak
into the numbers.

SSIM is implemented here rather than taken from skimage: skimage needs ~30 ms per
480x832 frame on CPU, and this evaluation compares ~500k frames.
"""

import torch
import torch.nn.functional as F


def _gaussian_window(size=11, sigma=1.5, device=None, dtype=torch.float32):
    coords = torch.arange(size, device=device, dtype=dtype) - (size - 1) / 2.0
    g = torch.exp(-(coords ** 2) / (2 * sigma ** 2))
    g = g / g.sum()
    return g[:, None] @ g[None, :]


class PixelMetrics:
    """Frame-wise metrics against a reference clip. Call once per (strategy, prompt)."""

    def __init__(self, device, lpips_net="alex", ssim_window=11, ssim_sigma=1.5):
        self.device = device
        self.window = _gaussian_window(ssim_window, ssim_sigma, device=device)[None, None]
        self.window_size = ssim_window
        import lpips
        self.lpips = lpips.LPIPS(net=lpips_net).to(device).eval()
        for p in self.lpips.parameters():
            p.requires_grad_(False)

    def ssim(self, x, y):
        """x, y: [T, C, H, W] in [0, 1]. Returns the mean over frames and channels."""
        t, c, h, w = x.shape
        win = self.window.expand(c, 1, self.window_size, self.window_size)
        pad = self.window_size // 2
        mu_x = F.conv2d(x, win, padding=pad, groups=c)
        mu_y = F.conv2d(y, win, padding=pad, groups=c)
        mu_x2, mu_y2, mu_xy = mu_x * mu_x, mu_y * mu_y, mu_x * mu_y
        sig_x = F.conv2d(x * x, win, padding=pad, groups=c) - mu_x2
        sig_y = F.conv2d(y * y, win, padding=pad, groups=c) - mu_y2
        sig_xy = F.conv2d(x * y, win, padding=pad, groups=c) - mu_xy
        c1, c2 = 0.01 ** 2, 0.03 ** 2
        s = ((2 * mu_xy + c1) * (2 * sig_xy + c2)) / ((mu_x2 + mu_y2 + c1) * (sig_x + sig_y + c2))
        # Drop the padded border, as skimage does with its default crop.
        s = s[..., pad:-pad, pad:-pad]
        return s.mean(dim=(1, 2, 3))  # per frame

    @torch.no_grad()
    def compute(self, frames, ref, chunk=16):
        """frames, ref: [T, C, H, W] float in [0, 1] on any device.

        Returns per-prompt mean MSE (for the protocol's ratio-of-means PSNR),
        plus mean SSIM and mean LPIPS over frames.
        """
        frames = frames.to(self.device, torch.float32)
        ref = ref.to(self.device, torch.float32)
        assert frames.shape == ref.shape, (frames.shape, ref.shape)

        mse_terms, ssim_terms, lpips_terms = [], [], []
        for i in range(0, frames.shape[0], chunk):
            a, b = frames[i:i + chunk], ref[i:i + chunk]
            mse_terms.append(((a - b) ** 2).mean(dim=(1, 2, 3)))
            ssim_terms.append(self.ssim(a, b))
            lpips_terms.append(self.lpips(a * 2 - 1, b * 2 - 1).flatten())

        mse = torch.cat(mse_terms)
        ssim = torch.cat(ssim_terms)
        lp = torch.cat(lpips_terms)
        mean_mse = mse.mean().item()
        return {
            "mean_mse": mean_mse,
            "psnr": -10.0 * torch.log10(torch.tensor(max(mean_mse, 1e-12))).item(),
            "ssim": ssim.mean().item(),
            "lpips": lp.mean().item(),
            "num_frames": int(frames.shape[0]),
        }

    @staticmethod
    def identity(num_frames):
        """Fixed self-comparison values the protocol pins for FFFF vs itself."""
        return {"mean_mse": 1e-12, "psnr": 120.0, "ssim": 1.0, "lpips": 0.0,
                "num_frames": int(num_frames)}