# Extended generation quality metrics for Phase 5 cross-family comparison.
# IS (Salimans et al., 2016): quality × diversity. LPIPS diversity: pairwise perceptual distance.
# Both functions accept float tensors in [-1, 1].
import torch
from torchmetrics.image.inception import InceptionScore
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity


# Inception Score (mean ± std) over 10 splits; imgs in [-1, 1], N ≥ 2048 for reliability; returns (mean, std)
def compute_is(
    imgs: torch.Tensor,
    device: str = "cuda",
    batch_size: int = 64,
) -> tuple[float, float]:
    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)


# Average pairwise LPIPS distance over n_pairs random (i ≠ j) pairs; higher = more diverse
def compute_lpips_diversity(
    imgs: torch.Tensor,
    n_pairs: int = 200,
    device: str = "cuda",
    batch_size: int = 16,
) -> float:
    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())
