tests/test_synth.py

"""WP-1: Tests for synthetic benchmark generation and evaluation metrics."""

from __future__ import annotations

import math
from pathlib import Path

import numpy as np
import pytest

from synth.generate import (
    generate_case,
    generate_dataset,
    save_case,
    load_case,
    load_fixtures,
    srgb_to_linear,
    linear_to_srgb,
)
from synth.evaluation import (
    psnr,
    ssim,
    degeneracy_indicator,
    score_pair,
    score_dataset,
    generate_report,
)


# ---------------------------------------------------------------------------
# Generation tests
# ---------------------------------------------------------------------------

class TestGenerateReproducible:
    """python -m synth.generate --n 6 must be reproducible by seed."""

    def test_same_seed_same_scan(self):
        c1 = generate_case(seed=0, stock="Generic", ratio=2.0, size=32)
        c2 = generate_case(seed=0, stock="Generic", ratio=2.0, size=32)
        assert np.allclose(c1["scan"], c2["scan"], atol=1e-6), \
            "Same seed produced different scans"

    def test_same_seed_same_gt(self):
        c1 = generate_case(seed=7, stock="Generic", ratio=3.0, size=32)
        c2 = generate_case(seed=7, stock="Generic", ratio=3.0, size=32)
        assert np.allclose(c1["gt_a"], c2["gt_a"], atol=1e-6)
        assert np.allclose(c1["gt_b"], c2["gt_b"], atol=1e-6)

    def test_different_seeds_different_scans(self):
        c1 = generate_case(seed=0, size=32)
        c2 = generate_case(seed=1, size=32)
        assert not np.allclose(c1["scan"], c2["scan"])

    def test_dataset_reproducible(self):
        ds1 = generate_dataset(n=4, seed=42, size=32)
        ds2 = generate_dataset(n=4, seed=42, size=32)
        for c1, c2 in zip(ds1, ds2):
            assert np.allclose(c1["scan"], c2["scan"], atol=1e-6)

    def test_dataset_n_cases(self):
        ds = generate_dataset(n=6, k1_fraction=0.25, size=32)
        assert len(ds) == 6
        n_k1 = sum(1 for c in ds if c["k1"])
        # round(6 * 0.25) = 2
        assert n_k1 == 2


class TestK1Cases:
    def test_k1_gt_b_is_black(self):
        c = generate_case(seed=0, k1=True, size=32)
        assert np.allclose(c["gt_b"], 0.0, atol=1e-6), "K=1 gt_b should be all zeros"

    def test_k1_h_b_is_zero(self):
        c = generate_case(seed=0, k1=True, size=32)
        assert c["h_b"].max() < 1e-8, "K=1 h_b should be zero"

    def test_k1_flag(self):
        c = generate_case(seed=0, k1=True, size=32)
        assert c["k1"] is True

    def test_k2_flag(self):
        c = generate_case(seed=0, k1=False, ratio=2.0, size=32)
        assert c["k1"] is False


class TestCaseStructure:
    def test_shapes(self):
        c = generate_case(seed=0, size=48)
        for key in ("scan", "gt_a", "gt_b"):
            assert c[key].shape == (48, 48, 3), f"{key} wrong shape: {c[key].shape}"
        for key in ("h_a", "h_b", "h_total"):
            assert c[key].shape == (48, 48), f"{key} wrong shape: {c[key].shape}"

    def test_scan_range(self):
        c = generate_case(seed=0, size=32)
        assert float(c["scan"].min()) >= 0.0
        assert float(c["scan"].max()) <= 1.0 + 1e-5

    def test_scan_mean_for_negative_appearance(self):
        """Scan should be bright (mean > 0.5) for the app's inversion heuristic."""
        for seed in range(3):
            c = generate_case(seed=seed, ratio=2.0, size=64)
            mean = float(c["scan"].mean())
            assert mean > 0.4, (
                f"Scan mean {mean:.3f} may be too low for app inversion heuristic "
                f"(seed={seed})"
            )

    def test_h_total_additivity(self):
        """h_total should equal h_a + h_b (before halation)."""
        c = generate_case(seed=0, halation=False, size=32)
        assert np.allclose(c["h_total"], c["h_a"] + c["h_b"], atol=1e-5)

    def test_halation_case_runs(self):
        c = generate_case(seed=0, halation=True, size=32)
        assert c["scan"].shape == (32, 32, 3)

    def test_jpeg_case_runs(self):
        c = generate_case(seed=0, jpeg=True, size=32)
        assert c["scan"].shape == (32, 32, 3)


class TestSaveLoad:
    def test_roundtrip(self, tmp_path):
        c = generate_case(seed=5, size=32)
        p = tmp_path / "case_000.npz"
        save_case(c, p)
        c2 = load_case(p)
        assert np.allclose(c["scan"], c2["scan"], atol=1e-6)
        assert np.allclose(c["gt_a"], c2["gt_a"], atol=1e-6)
        assert c["seed"] == c2["seed"]
        assert c["k1"] == c2["k1"]

    def test_load_fixtures(self):
        """Fixture cases must exist and load correctly."""
        cases = load_fixtures()
        assert len(cases) >= 6, f"Expected ≥6 fixture cases, got {len(cases)}"
        for c in cases:
            assert "scan" in c
            assert "gt_a" in c
            assert "gt_b" in c

    def test_load_fixtures_dir(self, tmp_path):
        """--fixtures-dir plumbing (WP-1.1): load_fixtures accepts custom dir (fast test, no 256px DIP)."""
        from synth.generate import generate_case, save_case
        for i in range(2):
            c = generate_case(seed=100 + i, size=32)
            p = tmp_path / f"case_{i:03d}.npz"
            save_case(c, p)
        cases = load_fixtures(fixtures_dir=tmp_path)
        assert len(cases) == 2
        assert all("scan" in c and "gt_a" in c for c in cases)

    def test_bench_limit_caps_tiny_dir(self, tmp_path):
        """--limit caps K2 selection in both benches (WP-1.2 fast test; tiny 64px dir, small budget, no full 256px run)."""
        import subprocess
        import sys
        from synth.generate import generate_case, save_case
        VENV_PY = sys.executable  # interpreter running pytest — portable, no hardcoded path
        # 3 guaranteed K=2 cases (k1=False)
        for i in range(3):
            c = generate_case(seed=300 + i, size=32, k1=False)
            p = tmp_path / f"case_{i:03d}.npz"
            save_case(c, p)
        # Test refine_bench --limit 1 (fast: 1 step)
        res1 = subprocess.run(
            [VENV_PY, "-m", "synth.refine_bench", "--bench", "--fixtures-dir", str(tmp_path),
             "--limit", "1", "--steps", "1"],
            capture_output=True, text=True, timeout=120
        )
        assert res1.returncode == 0, f"refine_bench failed: {res1.stderr}"
        assert "Using 1 K=2 fixtures" in res1.stdout
        # Test double_dip --limit 2 (fast: 5 iters)
        res2 = subprocess.run(
            [VENV_PY, "-m", "baselines.double_dip", "--bench", "--fixtures-dir", str(tmp_path),
             "--limit", "2", "--iters", "5"],
            capture_output=True, text=True, timeout=120
        )
        assert res2.returncode == 0, f"double_dip failed: {res2.stderr}"
        assert "Using 2 K=2 fixtures" in res2.stdout


# ---------------------------------------------------------------------------
# sRGB transfer function tests
# ---------------------------------------------------------------------------

class TestSRGBTransfer:
    def test_roundtrip(self):
        x = np.linspace(0.0, 1.0, 100).astype(np.float32)
        assert np.allclose(srgb_to_linear(linear_to_srgb(x)), x, atol=1e-5)
        assert np.allclose(linear_to_srgb(srgb_to_linear(x)), x, atol=1e-5)

    def test_midpoint(self):
        # sRGB 0.5 should linearize to approximately 0.214
        lin = float(srgb_to_linear(np.array([0.5]))[0])
        assert 0.20 < lin < 0.24, f"sRGB 0.5 → linear {lin:.4f}, expected ~0.214"


# ---------------------------------------------------------------------------
# Metric unit tests  (WP-1 acceptance criterion)
# ---------------------------------------------------------------------------

class TestPSNR:
    def test_identical_images_returns_inf(self):
        img = np.random.default_rng(0).uniform(0, 1, (32, 32, 3)).astype(np.float32)
        result = psnr(img, img)
        assert result == float("inf"), f"Expected inf, got {result}"

    def test_all_zeros_vs_all_ones(self):
        a = np.zeros((32, 32, 3), dtype=np.float32)
        b = np.ones((32, 32, 3), dtype=np.float32)
        result = psnr(a, b)
        assert result == 0.0  # MSE = 1.0, 20*log10(1/1) = 0

    def test_ordering(self):
        rng = np.random.default_rng(0)
        img = rng.uniform(0, 1, (32, 32, 3)).astype(np.float32)
        noise_small = img + 0.01 * rng.standard_normal(img.shape).astype(np.float32)
        noise_large = img + 0.1 * rng.standard_normal(img.shape).astype(np.float32)
        noise_small = np.clip(noise_small, 0, 1)
        noise_large = np.clip(noise_large, 0, 1)
        assert psnr(img, noise_small) > psnr(img, noise_large)


class TestSSIM:
    def test_identical_images(self):
        img = np.random.default_rng(0).uniform(0, 1, (32, 32, 3)).astype(np.float32)
        result = ssim(img, img)
        assert abs(result - 1.0) < 0.01, f"SSIM of identical images should be ~1, got {result}"

    def test_ordering(self):
        rng = np.random.default_rng(1)
        img = rng.uniform(0, 1, (32, 32, 3)).astype(np.float32)
        noise_small = np.clip(img + 0.01 * rng.standard_normal(img.shape), 0, 1).astype(np.float32)
        noise_large = np.clip(img + 0.2 * rng.standard_normal(img.shape), 0, 1).astype(np.float32)
        assert ssim(img, noise_small) > ssim(img, noise_large)


class TestDegeneracyIndicator:
    def test_equal_layers(self):
        img = np.full((32, 32, 3), 0.5, dtype=np.float32)
        result = degeneracy_indicator(img, img)
        assert abs(result - 0.5) < 0.01, f"Equal layers → degen ≈ 0.5, got {result}"

    def test_one_black_layer(self):
        img = np.full((32, 32, 3), 0.5, dtype=np.float32)
        black = np.zeros((32, 32, 3), dtype=np.float32)
        result = degeneracy_indicator(img, black)
        assert result < 0.05, f"One black layer → degen ≈ 0, got {result}"

    def test_both_black(self):
        black = np.zeros((32, 32, 3), dtype=np.float32)
        result = degeneracy_indicator(black, black)
        # Both black → undefined; implementation returns 0.5
        assert 0.0 <= result <= 1.0


class TestScorePair:
    """Acceptance criterion: identical images ⇒ PSNR=∞, swapped layers ⇒ same score."""

    def _make_images(self, seed=0):
        rng = np.random.default_rng(seed)
        gt_a = rng.uniform(0.1, 0.9, (32, 32, 3)).astype(np.float32)
        gt_b = rng.uniform(0.1, 0.9, (32, 32, 3)).astype(np.float32)
        pred_a = rng.uniform(0.1, 0.9, (32, 32, 3)).astype(np.float32)
        pred_b = rng.uniform(0.1, 0.9, (32, 32, 3)).astype(np.float32)
        return gt_a, gt_b, pred_a, pred_b

    def test_identical_images_psnr_inf(self):
        rng = np.random.default_rng(0)
        img_a = rng.uniform(0.1, 0.9, (32, 32, 3)).astype(np.float32)
        img_b = rng.uniform(0.1, 0.9, (32, 32, 3)).astype(np.float32)
        result = score_pair(img_a, img_b, img_a, img_b, compute_lpips=False)
        assert result["psnr"] == float("inf"), (
            f"Identical images should give PSNR=inf, got {result['psnr']}"
        )

    def test_swapped_layers_same_score(self):
        gt_a, gt_b, pred_a, pred_b = self._make_images()
        s1 = score_pair(gt_a, gt_b, pred_a, pred_b, compute_lpips=False)
        s2 = score_pair(gt_a, gt_b, pred_b, pred_a, compute_lpips=False)
        assert abs(s1["psnr"] - s2["psnr"]) < 1e-8, (
            f"Swapped preds gave different PSNR: {s1['psnr']:.6f} vs {s2['psnr']:.6f}"
        )
        assert abs(s1["ssim"] - s2["ssim"]) < 1e-8, (
            f"Swapped preds gave different SSIM: {s1['ssim']:.6f} vs {s2['ssim']:.6f}"
        )

    def test_returns_required_keys(self):
        gt_a, gt_b, pred_a, pred_b = self._make_images()
        result = score_pair(gt_a, gt_b, pred_a, pred_b, compute_lpips=False)
        for key in ("psnr", "ssim", "lpips", "density_mse", "degeneracy_indicator", "assignment"):
            assert key in result, f"Missing key: {key}"

    def test_better_prediction_higher_psnr(self):
        rng = np.random.default_rng(2)
        gt_a = rng.uniform(0.2, 0.8, (32, 32, 3)).astype(np.float32)
        gt_b = rng.uniform(0.2, 0.8, (32, 32, 3)).astype(np.float32)
        good_a = np.clip(gt_a + 0.01 * rng.standard_normal(gt_a.shape), 0, 1).astype(np.float32)
        good_b = np.clip(gt_b + 0.01 * rng.standard_normal(gt_b.shape), 0, 1).astype(np.float32)
        bad_a = rng.uniform(0.1, 0.9, (32, 32, 3)).astype(np.float32)
        bad_b = rng.uniform(0.1, 0.9, (32, 32, 3)).astype(np.float32)
        good_score = score_pair(gt_a, gt_b, good_a, good_b, compute_lpips=False)
        bad_score = score_pair(gt_a, gt_b, bad_a, bad_b, compute_lpips=False)
        assert good_score["psnr"] > bad_score["psnr"]


# ---------------------------------------------------------------------------
# Integration: ranking demo candidates on fixtures produces a report
# ---------------------------------------------------------------------------

class TestRankingProducesReport:
    def test_benchmark_report_artifact(self, tmp_path):
        """
        Rank existing demo candidates on the fixture set → produces a Markdown
        report artifact. Verifies the measurement pipeline end-to-end.
        """
        from PIL import Image as PILImage
        from app.preprocessing import preprocess_negative, to_pil
        from app.api_client import generate_demo_candidates
        from app.scoring import rank_candidates
        from film_physics import get_film_curve

        cases = load_fixtures()
        results = []

        for case in cases:
            # Convert float32 [0,1] scan to PIL (the app's input format)
            scan_uint8 = (case["scan"] * 255.0).clip(0, 255).astype(np.uint8)
            scan_pil = PILImage.fromarray(scan_uint8)

            preprocessed = preprocess_negative(scan_pil)

            candidates = generate_demo_candidates(preprocessed.rgb, num_candidates=2)

            film_curve = get_film_curve(case["stock"])
            ranked = rank_candidates(
                candidates=candidates,
                observed_log_exposure=preprocessed.log_exposure,
                observed_rgb=preprocessed.rgb,
                film_curve=film_curve,
                density=preprocessed.density,
                confidence_mask=preprocessed.confidence_mask,
            )
            best = ranked[0]

            score = score_pair(
                gt_a=case["gt_a"],
                gt_b=case["gt_b"],
                pred_a=best.separation.image_a,
                pred_b=best.separation.image_b,
                compute_lpips=False,
            )
            score["seed"] = case["seed"]
            score["ratio"] = case["ratio"]
            score["stock"] = case["stock"]
            score["k1"] = case["k1"]
            results.append(score)

        report_path = tmp_path / "benchmark_report.md"
        report_md = generate_report(results, output_path=str(report_path))

        assert report_path.exists(), "Report file was not created"
        assert len(report_md) > 100, "Report is suspiciously short"
        assert "PSNR" in report_md
        assert "SSIM" in report_md
        # JSON sidecar
        json_path = report_path.with_suffix(".json")
        assert json_path.exists(), "JSON sidecar was not created"


# ---------------------------------------------------------------------------
# WP-3 required tests (Part V)
# ---------------------------------------------------------------------------

from app.api_client import SeparationResult


class TestWP3DegenerateAndKSelection:
    """Required acceptance from WP-3 spec."""

    def test_degenerate_ranks_worse_than_nondeg_on_ratio_le4(self):
        """
        (a) Explicit degenerate (A=observed positive, B=black) must rank
        strictly worse (higher loss) than every non-degenerate demo candidate
        on all fixtures with exposure ratio <=4:1 .
        """
        from PIL import Image as PILImage
        from app.preprocessing import preprocess_negative
        from app.api_client import generate_demo_candidates
        from app.scoring import rank_candidates, score_separation
        from film_physics import get_film_curve
        from densitometry import VALID

        cases = load_fixtures()

        # (Fix 5 / review #7) on-the-fly 3.5:1 case + degen balance_loss >=0.5  -- runs ONCE
        from synth.generate import generate_case
        case35 = generate_case(seed=123, ratio=3.5, size=128, stock="Generic")
        pil35 = PILImage.fromarray((case35["scan"] * 255).clip(0, 255).astype(np.uint8))
        pre35 = preprocess_negative(pil35, stock="Generic")
        cands35 = generate_demo_candidates(pre35.rgb, num_candidates=2)
        ranked35 = rank_candidates(
            cands35, pre35.log_exposure, pre35.rgb, get_film_curve("Generic"),
            density=pre35.density, confidence_mask=pre35.confidence_mask,
        )
        best35 = ranked35[0]
        degen35_bd = score_separation(
            pre35.log_exposure, pre35.rgb, pre35.rgb, np.zeros_like(pre35.rgb),
            get_film_curve("Generic"), density=pre35.density, confidence_mask=pre35.confidence_mask,
        )
        margin = best35.score.total_loss - degen35_bd.total_loss
        assert margin < -0.1, f"Expected degen worse by >0.1 total, margin={margin}"
        assert getattr(degen35_bd, "balance_loss", 0.0) >= 0.5, f"degen balance_loss must be >=0.5, got {getattr(degen35_bd, 'balance_loss', 0)}"

        for case in cases:
            r = float(case["ratio"])
            if case["k1"] or r > 4.0 + 1e-6:
                continue
            scan_uint = (case["scan"] * 255).clip(0, 255).astype(np.uint8)
            pil = PILImage.fromarray(scan_uint)
            pre = preprocess_negative(pil, stock=case["stock"])

            demo_cands = generate_demo_candidates(pre.rgb, num_candidates=4)
            film_curve = get_film_curve(case["stock"])

            # score non-deg
            nondeg_losses = []
            for c in demo_cands:
                bd = score_separation(
                    observed_log_exposure=pre.log_exposure,
                    observed_rgb=pre.rgb,
                    image_a_rgb=c.image_a,
                    image_b_rgb=c.image_b,
                    film_curve=film_curve,
                    density=pre.density,
                    confidence_mask=pre.confidence_mask,
                )
                nondeg_losses.append(bd.total_loss)

            # explicit degenerate
            black = np.zeros_like(pre.rgb)
            degen = SeparationResult(
                image_a=pre.rgb.copy(),
                image_b=black,
                method="degenerate",
                message="test degen",
                candidate_id="degen",
            )
            degen_bd = score_separation(
                observed_log_exposure=pre.log_exposure,
                observed_rgb=pre.rgb,
                image_a_rgb=degen.image_a,
                image_b_rgb=degen.image_b,
                film_curve=film_curve,
                density=pre.density,
                confidence_mask=pre.confidence_mask,
            )

            for ndl in nondeg_losses:
                assert degen_bd.total_loss > ndl + 1e-9, (
                    f"Degen loss {degen_bd.total_loss:.6f} not > nondeg {ndl:.6f} "
                    f"on fixture r={r:.2f}"
                )

    def test_k_selection_flags_k1_and_lowratio_k2(self):
        """
        (b) K-selection score must flag K=1 fixtures as single-exposure (low score)
        and ratio <=2:1 fixtures as K=2 (high score).
        """
        from PIL import Image as PILImage
        from app.preprocessing import preprocess_negative
        from app.api_client import generate_demo_candidates
        from app.scoring import rank_candidates
        from film_physics import get_film_curve

        cases = load_fixtures()
        for case in cases:
            scan_uint = (case["scan"] * 255).clip(0, 255).astype(np.uint8)
            pil = PILImage.fromarray(scan_uint)
            pre = preprocess_negative(pil, stock=case["stock"])
            cands = generate_demo_candidates(pre.rgb, num_candidates=2)
            film_curve = get_film_curve(case["stock"])
            ranked = rank_candidates(
                candidates=cands,
                observed_log_exposure=pre.log_exposure,
                observed_rgb=pre.rgb,
                film_curve=film_curve,
                density=pre.density,
                confidence_mask=pre.confidence_mask,
            )
            # Use the top ranked's k score as the frame judgment
            ksel = ranked[0].score.k_selection_score if hasattr(ranked[0].score, "k_selection_score") else 0.0
            assert 0.0 <= ksel <= 1.0
            is_k1 = bool(case["k1"])
            r = float(case["ratio"])
            if is_k1:
                assert ksel < 0.5, f"K=1 fixture got ksel={ksel:.2f} (should flag single)"
            elif r <= 2.0 + 1e-6:
                assert ksel >= 0.35, f"ratio<={r:.1f} fixture got ksel={ksel:.2f} (should flag K=2)"