"""Extended generation quality metrics for Phase 5 cross-family comparison.

IS  — Inception Score (Salimans et al., 2016): measures sample quality × diversity.
LPIPS — average pairwise learned perceptual distance: measures sample diversity alone.

Both functions accept float tensors in [-1, 1].
"""
import torch
from torchmetrics.image.inception import InceptionScore
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity


def compute_is(
    imgs: torch.Tensor,
    device: str = "cuda",
    batch_size: int = 64,
) -> tuple[float, float]:
    """Inception Score (mean ± std) over 10 splits.

    imgs: (N, 3, H, W) in [-1, 1].  N ≥ 2 048 for a reliable estimate.
    Returns (is_mean, is_std).
    """
    metric = InceptionScore(normalize=True).to(device)
    imgs_01 = (imgs.clamp(-1, 1) * 0.5 + 0.5)
    for i in range(0, len(imgs_01), batch_size):
        metric.update(imgs_01[i : i + batch_size].to(device))
    mean, std = metric.compute()
    return float(mean), float(std)


def compute_lpips_diversity(
    imgs: torch.Tensor,
    n_pairs: int = 200,
    device: str = "cuda",
    batch_size: int = 16,
) -> float:
    """Average pairwise LPIPS distance — higher means more diverse samples.

    imgs: (N, 3, H, W) in [-1, 1].  Samples n_pairs random (i, j) pairs with i ≠ j.
    """
    metric = LearnedPerceptualImagePatchSimilarity(net_type="alex", normalize=True).to(device)
    imgs_01 = (imgs.clamp(-1, 1) * 0.5 + 0.5)
    N = len(imgs_01)

    # Sample random pairs (ensure i ≠ j by rejection)
    idx = torch.randperm(N * 2)[:n_pairs * 2].view(n_pairs, 2) % N
    same = idx[:, 0] == idx[:, 1]
    idx[same, 1] = (idx[same, 1] + 1) % N  # shift duplicate indices

    for start in range(0, n_pairs, batch_size):
        end = min(start + batch_size, n_pairs)
        i_batch = idx[start:end, 0]
        j_batch = idx[start:end, 1]
        metric.update(imgs_01[i_batch].to(device), imgs_01[j_batch].to(device))

    return float(metric.compute())
