synth/evaluation.py

"""
Evaluation metrics and report generation for the double-exposure benchmark (WP-1).

All metrics are permutation-invariant: both (pred_a→gt_a, pred_b→gt_b) and
(pred_a→gt_b, pred_b→gt_a) assignments are scored; the better assignment
(lower total LPIPS) is reported.

PSNR=∞ guard: when MSE == 0 (identical images), returns float('inf').
"""

from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

# Lazy-import heavy dependencies
_SKIMAGE_AVAILABLE = None
_LPIPS_CACHE: dict = {}


def _check_skimage() -> bool:
    global _SKIMAGE_AVAILABLE
    if _SKIMAGE_AVAILABLE is None:
        try:
            import skimage  # noqa: F401
            _SKIMAGE_AVAILABLE = True
        except ImportError:
            _SKIMAGE_AVAILABLE = False
    return _SKIMAGE_AVAILABLE


# ---------------------------------------------------------------------------
# Per-image metrics
# ---------------------------------------------------------------------------

def psnr(img1: np.ndarray, img2: np.ndarray) -> float:
    """Peak signal-to-noise ratio (dB). Returns float('inf') for identical images."""
    mse = float(np.mean((img1.astype(np.float64) - img2.astype(np.float64)) ** 2))
    if mse == 0.0:
        return float("inf")
    return float(20.0 * math.log10(1.0 / math.sqrt(mse)))


def ssim(img1: np.ndarray, img2: np.ndarray) -> float:
    """
    Structural Similarity Index (SSIM) via scikit-image.

    Falls back to a basic luminance-correlation estimate if scikit-image is
    unavailable (noted in result via the 'ssim_approx' key).
    """
    if _check_skimage():
        from skimage.metrics import structural_similarity
        # Handle grayscale and RGB
        channel_axis = -1 if img1.ndim == 3 else None
        return float(
            structural_similarity(
                img1.astype(np.float64),
                img2.astype(np.float64),
                data_range=1.0,
                channel_axis=channel_axis,
            )
        )
    # Fallback: normalized cross-correlation (rough approximation)
    mu1, mu2 = img1.mean(), img2.mean()
    s1, s2 = img1.std(), img2.std()
    cov = float(np.mean((img1 - mu1) * (img2 - mu2)))
    denom = (s1 * s2) + 1e-8
    return float(cov / denom)


def _get_lpips(net: str = "alex"):
    if net not in _LPIPS_CACHE:
        import lpips
        import torch
        model = lpips.LPIPS(net=net)
        model.eval()
        for p in model.parameters():
            p.requires_grad = False
        _LPIPS_CACHE[net] = model
    return _LPIPS_CACHE[net]


def lpips_distance(img1: np.ndarray, img2: np.ndarray, net: str = "alex") -> float:
    """
    LPIPS perceptual distance between two (H, W, 3) float32 images in [0, 1].

    Returns NaN if LPIPS cannot be computed (import error).
    """
    try:
        import torch
        model = _get_lpips(net)

        def to_t(img: np.ndarray):
            t = torch.from_numpy(img).float().permute(2, 0, 1).unsqueeze(0)
            return t * 2.0 - 1.0  # [0,1] → [-1,1]

        with torch.no_grad():
            dist = model(to_t(img1), to_t(img2)).mean()
        return float(dist.item())
    except Exception:
        return float("nan")


# ---------------------------------------------------------------------------
# Degeneracy indicator
# ---------------------------------------------------------------------------

def degeneracy_indicator(pred_a: np.ndarray, pred_b: np.ndarray) -> float:
    """
    Minimum layer energy share — close to 0 means one layer is near-black (degenerate).

    Computes mean luminance of each predicted layer and returns
    min(share_a, share_b) where share_a = mean_a / (mean_a + mean_b).
    """
    def lum(img: np.ndarray) -> float:
        g = img.mean(axis=-1) if img.ndim == 3 else img
        return float(g.mean())

    la = lum(pred_a)
    lb = lum(pred_b)
    total = la + lb
    if total < 1e-10:
        return 0.5  # both black — undefined; return balanced
    share_a = la / total
    return float(min(share_a, 1.0 - share_a))


# ---------------------------------------------------------------------------
# Density residual
# ---------------------------------------------------------------------------

def density_residual_mse(
    pred_a: np.ndarray,
    pred_b: np.ndarray,
    h_total: np.ndarray,
    film_curve,
    valid_mask: Optional[np.ndarray] = None,
) -> float:
    """
    Density-space recombination fidelity on the valid mask.

    Given predicted layers pred_a, pred_b (sRGB positive), computes:
      1. Linearize and extract luminance → H_pred_a, H_pred_b (up to a scale).
      2. Find optimal global scale g (least squares) so g*(H_pred_a + H_pred_b) ≈ H_total.
      3. Compute density of each via the forward curve.
      4. Return MSE between predicted and GT density on valid_mask.

    Returns NaN if computation fails.
    """
    try:
        from synth.generate import srgb_to_linear
        import torch

        def lum_from_srgb(img: np.ndarray) -> np.ndarray:
            lin = srgb_to_linear(img)
            return (0.2126 * lin[..., 0] + 0.7152 * lin[..., 1] + 0.0722 * lin[..., 2]).astype(np.float32)

        y_a = lum_from_srgb(pred_a)
        y_b = lum_from_srgb(pred_b)
        y_sum = y_a + y_b  # unnormalized

        # Optimal scale g: minimize ||g*y_sum - h_total||^2
        numer = float(np.sum(y_sum * h_total))
        denom = float(np.sum(y_sum ** 2)) + 1e-10
        g = max(numer / denom, 1e-6)

        h_pred_total = np.clip(g * y_sum, 1e-8, None)

        # Build valid mask: mid-range of H_total (not toe-noise, not shoulder-saturated)
        if valid_mask is None:
            h_norm = h_total / (h_total.max() + 1e-8)
            valid_mask = (h_norm > 0.05) & (h_norm < 0.95)

        if not valid_mask.any():
            return float("nan")

        # Apply forward curve to both
        def apply_curve(h: np.ndarray) -> np.ndarray:
            log_h = np.log10(np.clip(h, 1e-8, None))
            t = torch.from_numpy(log_h).float().unsqueeze(0).unsqueeze(0)
            with torch.no_grad():
                d = film_curve(t)
            return d.squeeze().numpy().astype(np.float32)

        d_pred = apply_curve(h_pred_total)
        d_gt = apply_curve(h_total)

        mse = float(np.mean((d_pred[valid_mask] - d_gt[valid_mask]) ** 2))
        return mse
    except Exception:
        return float("nan")


# ---------------------------------------------------------------------------
# Permutation-invariant pair scoring
# ---------------------------------------------------------------------------

def _score_assignment(
    gt_x: np.ndarray,
    gt_y: np.ndarray,
    pred_a: np.ndarray,
    pred_b: np.ndarray,
    compute_lpips: bool,
    film_curve,
    h_total: Optional[np.ndarray],
) -> Dict[str, float]:
    """Score pred_a→gt_x and pred_b→gt_y."""
    p_a = psnr(pred_a, gt_x)
    p_b = psnr(pred_b, gt_y)
    s_a = ssim(pred_a, gt_x)
    s_b = ssim(pred_b, gt_y)

    if compute_lpips:
        l_a = lpips_distance(pred_a, gt_x)
        l_b = lpips_distance(pred_b, gt_y)
    else:
        l_a = l_b = float("nan")

    dm = float("nan")
    if film_curve is not None and h_total is not None:
        dm = density_residual_mse(pred_a, pred_b, h_total, film_curve)

    degen = degeneracy_indicator(pred_a, pred_b)

    # Gain-matched variant (important for high-ratio weak layer which is dark but may be correctly recovered at different scale)
    try:
        p_a_g = _gain_matched_psnr(gt_x, pred_a)
        p_b_g = _gain_matched_psnr(gt_y, pred_b)
    except Exception:
        p_a_g = p_a
        p_b_g = p_b

    return {
        "psnr_a": p_a,
        "psnr_b": p_b,
        "psnr": _mean_finite(p_a, p_b),
        "psnr_gain_matched": _mean_finite(p_a_g, p_b_g),
        "ssim_a": s_a,
        "ssim_b": s_b,
        "ssim": _mean_finite(s_a, s_b),
        "lpips_a": l_a,
        "lpips_b": l_b,
        "lpips": _mean_finite(l_a, l_b),
        "density_mse": dm,
        "degeneracy_indicator": degen,
        "assignment": "ab",
    }


def _mean_finite(*vals) -> float:
    """Mean of values, excluding only NaN. inf is kept so PSNR=∞ propagates correctly."""
    valid = [v for v in vals if not math.isnan(v)]
    return sum(valid) / len(valid) if valid else float("nan")


def _gain_matched_psnr(gt: np.ndarray, pred: np.ndarray) -> float:
    """Fit scalar gain g to minimize ||gt - g*pred|| then return PSNR on the matched pair."""
    g = np.dot(gt.ravel().astype(np.float64), pred.ravel().astype(np.float64)) / (np.dot(pred.ravel().astype(np.float64), pred.ravel().astype(np.float64)) + 1e-12)
    g = max(g, 1e-6)
    matched = np.clip(g * pred, 0.0, 1.0)
    return psnr(gt, matched)


def score_pair(
    gt_a: np.ndarray,
    gt_b: np.ndarray,
    pred_a: np.ndarray,
    pred_b: np.ndarray,
    film_curve=None,
    h_total: Optional[np.ndarray] = None,
    compute_lpips: bool = True,
) -> Dict[str, float]:
    """
    Permutation-invariant scoring of a recovered pair against ground truth.

    Both assignments (pred_a→gt_a, pred_b→gt_b) and (pred_a→gt_b, pred_b→gt_a)
    are evaluated; the assignment with lower mean LPIPS (or lower mean PSNR
    difference when LPIPS is unavailable) is returned.

    PSNR=∞ guard: identical images → psnr = float('inf').
    Swapped-layers invariance: score_pair(gt_a, gt_b, pred_a, pred_b) ==
                               score_pair(gt_a, gt_b, pred_b, pred_a).
    """
    s_ab = _score_assignment(gt_a, gt_b, pred_a, pred_b, compute_lpips, film_curve, h_total)
    s_ba = _score_assignment(gt_b, gt_a, pred_a, pred_b, compute_lpips, film_curve, h_total)
    s_ba["assignment"] = "ba"

    # Select better assignment: lower LPIPS if finite, else higher PSNR
    lpips_ab = _mean_finite(s_ab["lpips_a"], s_ab["lpips_b"])
    lpips_ba = _mean_finite(s_ba["lpips_a"], s_ba["lpips_b"])

    if math.isfinite(lpips_ab) and math.isfinite(lpips_ba):
        return s_ab if lpips_ab <= lpips_ba else s_ba
    else:
        # Fall back to PSNR (higher is better)
        psnr_ab = _mean_finite(s_ab["psnr_a"], s_ab["psnr_b"])
        psnr_ba = _mean_finite(s_ba["psnr_a"], s_ba["psnr_b"])
        return s_ab if psnr_ab >= psnr_ba else s_ba


# ---------------------------------------------------------------------------
# Dataset-level scoring
# ---------------------------------------------------------------------------

def score_dataset(
    cases: List[dict],
    predictions: List[Tuple[np.ndarray, np.ndarray]],
    compute_lpips: bool = True,
) -> List[Dict[str, Any]]:
    """
    Score all (case, prediction) pairs.

    Args:
        cases: List of case dicts (from generate_dataset or load_fixtures).
        predictions: List of (pred_a, pred_b) pairs aligned with cases.
        compute_lpips: Whether to compute (slow) LPIPS.

    Returns:
        List of per-case result dicts.
    """
    from film_physics import get_film_curve

    results = []
    for case, (pred_a, pred_b) in zip(cases, predictions):
        score = score_pair(
            gt_a=case["gt_a"],
            gt_b=case["gt_b"],
            pred_a=pred_a,
            pred_b=pred_b,
            film_curve=get_film_curve(case["stock"]),
            h_total=case["h_total"],
            compute_lpips=compute_lpips,
        )
        score["seed"] = case["seed"]
        score["ratio"] = case["ratio"]
        score["stock"] = case["stock"]
        score["k1"] = case["k1"]
        results.append(score)
    return results


# ---------------------------------------------------------------------------
# Report generation
# ---------------------------------------------------------------------------

def _ratio_band(ratio: float, k1: bool) -> str:
    if k1:
        return "K=1 (control)"
    if ratio < 2.0:
        return "1:1 – 2:1"
    if ratio < 4.0:
        return "2:1 – 4:1"
    return "4:1 – 8:1+"


def _fmt(val: float, fmt: str = ".4f") -> str:
    if not math.isfinite(val):
        return "∞" if val == float("inf") else "NaN"
    return format(val, fmt)


def generate_report(
    results: List[Dict[str, Any]],
    output_path: Optional[str] = None,
) -> str:
    """
    Generate a Markdown + JSON benchmark report.

    Args:
        results: Per-case result dicts from score_dataset or score_pair calls.
        output_path: If given, writes the .md file and a .json sidecar.

    Returns:
        The Markdown report string.
    """
    if not results:
        return "# Benchmark Report\n\nNo results.\n"

    def safe_mean(vals: List[float]) -> float:
        finite = [v for v in vals if math.isfinite(v)]
        return sum(finite) / len(finite) if finite else float("nan")

    # Overall summary
    psnrs = [r["psnr"] for r in results]
    psnrs_g = [r.get("psnr_gain_matched", r["psnr"]) for r in results]
    ssims = [r["ssim"] for r in results]
    lpipss = [r["lpips"] for r in results]
    dmses = [r["density_mse"] for r in results]
    degens = [r["degeneracy_indicator"] for r in results]

    lines: List[str] = [
        "# Double-Exposure Benchmark Report",
        "",
        "## Summary",
        "",
        "| Metric | Mean |",
        "|--------|------|",
        f"| PSNR (dB) | {_fmt(safe_mean(psnrs))} |",
        f"| PSNR (gain-matched) | {_fmt(safe_mean(psnrs_g))} |",
        f"| SSIM | {_fmt(safe_mean(ssims))} |",
        f"| LPIPS | {_fmt(safe_mean(lpipss))} |",
        f"| Density MSE | {_fmt(safe_mean(dmses))} |",
        f"| Degeneracy indicator | {_fmt(safe_mean(degens))} |",
        "",
        "## Stratified by Exposure Ratio",
        "",
        "| Ratio band | N | PSNR | SSIM | LPIPS | Degeneracy |",
        "|------------|---|------|------|-------|------------|",
    ]

    bands: Dict[str, List[dict]] = {}
    for r in results:
        band = _ratio_band(r["ratio"], r["k1"])
        bands.setdefault(band, []).append(r)

    band_order = ["1:1 – 2:1", "2:1 – 4:1", "4:1 – 8:1+", "K=1 (control)"]
    for band in band_order:
        if band not in bands:
            continue
        rs = bands[band]
        n = len(rs)
        lines.append(
            f"| {band} | {n} "
            f"| {_fmt(safe_mean([r['psnr'] for r in rs]))} "
            f"| {_fmt(safe_mean([r['ssim'] for r in rs]))} "
            f"| {_fmt(safe_mean([r['lpips'] for r in rs]))} "
            f"| {_fmt(safe_mean([r['degeneracy_indicator'] for r in rs]))} |"
        )

    lines += [
        "",
        "## Per-Case Details",
        "",
        "| Seed | K | Ratio | Stock | PSNR | SSIM | LPIPS | Degen. |",
        "|------|---|-------|-------|------|------|-------|--------|",
    ]
    for r in results:
        k_label = "1" if r["k1"] else "2"
        ratio_str = "∞" if r["ratio"] > 100 else f"{r['ratio']:.2f}"
        lines.append(
            f"| {r['seed']} | {k_label} | {ratio_str} | {r['stock']} "
            f"| {_fmt(r['psnr'])} "
            f"| {_fmt(r['ssim'])} "
            f"| {_fmt(r['lpips'])} "
            f"| {_fmt(r['degeneracy_indicator'])} |"
        )

    report_md = "\n".join(lines) + "\n"

    if output_path is not None:
        out = Path(output_path)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(report_md, encoding="utf-8")
        # JSON sidecar
        json_path = out.with_suffix(".json")
        json_path.write_text(
            json.dumps(
                {
                    "summary": {
                        "psnr": safe_mean(psnrs),
                        "ssim": safe_mean(ssims),
                        "lpips": safe_mean(lpipss),
                        "density_mse": safe_mean(dmses),
                        "degeneracy_indicator": safe_mean(degens),
                    },
                    "per_case": results,
                },
                indent=2,
                default=lambda x: None if (isinstance(x, float) and not math.isfinite(x)) else x,
            ),
            encoding="utf-8",
        )

    return report_md