tests/test_densitometry.py

"""
Tests for WP-2: densitometry module.

Covers:
  - sRGB round-trip identity
  - PiecewiseFilmCurve.inverse() round-trip error < 0.01 D
  - scan_to_density smoke test
  - density_to_h_total: confidence mask values
  - Fixture validation: Ĥ_total matches GT h_total within 5% median relative
    error on the VALID mask, up to one global scale
"""

from __future__ import annotations

import math
from pathlib import Path

import numpy as np
import pytest

from densitometry import (
    TOE,
    VALID,
    SHOULDER,
    srgb_to_linear,
    linear_to_srgb,
    scan_to_density,
    density_to_h_total,
    scan_to_density_rgb,
    density_to_h_total_rgb,
    combine_confidence_rgb,
)
from film_physics import get_film_curve


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

FIXTURES_DIR = Path(__file__).parent.parent / "synth" / "fixtures"


@pytest.fixture(scope="module")
def generic_curve():
    return get_film_curve("Generic")


@pytest.fixture(scope="module")
def all_fixtures():
    paths = sorted(FIXTURES_DIR.glob("case_*.npz"))
    if not paths:
        pytest.skip("No fixture cases found — run WP-1 generator first")
    cases = []
    for p in paths:
        d = np.load(p, allow_pickle=False)
        cases.append({
            "scan": d["scan"],
            "h_total": d["h_total"],
            "stock": str(d["stock"]),
        })
    return cases


# ---------------------------------------------------------------------------
# sRGB round-trip
# ---------------------------------------------------------------------------

class TestSRGBRoundtrip:
    def test_linear_to_srgb_to_linear(self):
        rng = np.random.default_rng(0)
        x = rng.uniform(0.0, 1.0, (64, 64, 3)).astype(np.float32)
        recovered = srgb_to_linear(linear_to_srgb(x))
        np.testing.assert_allclose(recovered, x, atol=1e-5)

    def test_srgb_to_linear_to_srgb(self):
        rng = np.random.default_rng(1)
        x = rng.uniform(0.0, 1.0, (64, 64, 3)).astype(np.float32)
        recovered = linear_to_srgb(srgb_to_linear(x))
        np.testing.assert_allclose(recovered, x, atol=1e-5)

    def test_boundary_values(self):
        # Endpoints must be exact
        x = np.array([[[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]], dtype=np.float32)
        lin = srgb_to_linear(x)
        assert float(lin[0, 0, 0]) == pytest.approx(0.0, abs=1e-7)
        assert float(lin[0, 1, 0]) == pytest.approx(1.0, abs=1e-7)


# ---------------------------------------------------------------------------
# PiecewiseFilmCurve.inverse() round-trip
# ---------------------------------------------------------------------------

class TestCurveInverse:
    def test_roundtrip_below_threshold(self, generic_curve):
        d_min = float(generic_curve.d_min)
        d_max = float(generic_curve.d_max)
        # Sample D values strictly inside the valid range
        d_test = np.linspace(d_min + 0.05, d_max - 0.05, 500).astype(np.float32)
        log_h = generic_curve.inverse(d_test)
        import torch
        with torch.no_grad():
            d_recovered = generic_curve.forward(torch.from_numpy(log_h)).numpy()
        max_err = float(np.max(np.abs(d_recovered - d_test)))
        assert max_err < 0.01, f"Max round-trip error {max_err:.5f} D ≥ 0.01"

    def test_inverse_monotone(self, generic_curve):
        d_min = float(generic_curve.d_min)
        d_max = float(generic_curve.d_max)
        d_vals = np.linspace(d_min + 0.01, d_max - 0.01, 200)
        log_h = generic_curve.inverse(d_vals)
        diffs = np.diff(log_h)
        assert np.all(diffs >= -1e-6), "inverse() must be monotone non-decreasing"

    def test_all_stocks(self):
        from film_physics import list_film_stocks
        for stock in list_film_stocks():
            curve = get_film_curve(stock)
            d_min = float(curve.d_min)
            d_max = float(curve.d_max)
            d_test = np.linspace(d_min + 0.05, d_max - 0.05, 200).astype(np.float32)
            log_h = curve.inverse(d_test)
            import torch
            with torch.no_grad():
                d_back = curve.forward(torch.from_numpy(log_h)).numpy()
            err = float(np.max(np.abs(d_back - d_test)))
            assert err < 0.01, f"{stock}: round-trip error {err:.5f} D ≥ 0.01"

    def test_2d_input_shape(self, generic_curve):
        d_2d = np.full((8, 8), 0.5, dtype=np.float32)
        log_h = generic_curve.inverse(d_2d)
        assert log_h.shape == (8, 8)


# ---------------------------------------------------------------------------
# scan_to_density smoke tests
# ---------------------------------------------------------------------------

class TestScanToDensity:
    def test_output_shapes(self, generic_curve):
        rng = np.random.default_rng(42)
        scan = rng.uniform(0.4, 0.9, (32, 32, 3)).astype(np.float32)
        d, lin = scan_to_density(scan, stock="Generic")
        assert d.shape == (32, 32)
        assert lin.shape == (32, 32, 3)

    def test_density_nonnegative(self):
        rng = np.random.default_rng(7)
        scan = rng.uniform(0.3, 1.0, (32, 32, 3)).astype(np.float32)
        d, _ = scan_to_density(scan)
        assert float(d.min()) >= -1e-4, "Density should be non-negative"

    def test_white_level_override(self, generic_curve):
        rng = np.random.default_rng(3)
        scan = rng.uniform(0.5, 0.9, (32, 32, 3)).astype(np.float32)
        d1, _ = scan_to_density(scan)
        d2, _ = scan_to_density(scan, white_level=float(np.percentile(scan.mean(-1), 99.5)))
        # Both should be finite
        assert np.all(np.isfinite(d1))
        assert np.all(np.isfinite(d2))

    def test_synthetic_roundtrip(self, generic_curve):
        """Forward-model scan → scan_to_density should recover D close to original."""
        from densitometry import linear_to_srgb as l2s, srgb_to_linear as s2l
        rng = np.random.default_rng(99)
        # Build a simple test: constant H_total at midtone
        h = np.full((32, 32), 0.4, dtype=np.float32)
        import torch
        log_h = np.log10(h)
        with torch.no_grad():
            d_gt = generic_curve.forward(torch.from_numpy(log_h)).numpy()
        t = np.power(10.0, -d_gt)
        scan_gray = l2s(t)
        scan_rgb = np.stack([scan_gray] * 3, axis=-1)

        d_rec, _ = scan_to_density(scan_rgb, stock="Generic")
        # Allow for white-point estimation uncertainty (< 0.05 D)
        err = float(np.median(np.abs(d_rec - d_gt)))
        assert err < 0.05, f"Median density recovery error {err:.4f} ≥ 0.05"


# ---------------------------------------------------------------------------
# density_to_h_total — confidence mask
# ---------------------------------------------------------------------------

class TestDensityToHTotalMask:
    def test_mask_values_valid(self, generic_curve):
        d_min = float(generic_curve.d_min)
        d_max = float(generic_curve.d_max)
        mid = (d_min + d_max) / 2.0
        d_arr = np.full((10, 10), mid, dtype=np.float32)
        _, mask = density_to_h_total(d_arr, generic_curve)
        assert np.all(mask == VALID)

    def test_mask_values_toe(self, generic_curve):
        d_min = float(generic_curve.d_min)
        d_arr = np.full((10, 10), d_min, dtype=np.float32)
        _, mask = density_to_h_total(d_arr, generic_curve)
        assert np.all(mask == TOE)

    def test_mask_values_shoulder(self, generic_curve):
        d_max = float(generic_curve.d_max)
        d_arr = np.full((10, 10), d_max, dtype=np.float32)
        _, mask = density_to_h_total(d_arr, generic_curve)
        assert np.all(mask == SHOULDER)

    def test_h_total_positive(self, generic_curve):
        d_min = float(generic_curve.d_min)
        d_max = float(generic_curve.d_max)
        d_arr = np.linspace(d_min, d_max, 100).reshape(10, 10).astype(np.float32)
        h, _ = density_to_h_total(d_arr, generic_curve)
        assert np.all(h > 0), "H_total must be strictly positive"

    def test_h_total_monotone_in_density(self, generic_curve):
        d_min = float(generic_curve.d_min)
        d_max = float(generic_curve.d_max)
        d_arr = np.linspace(d_min + 0.01, d_max - 0.01, 100).astype(np.float32)
        h, _ = density_to_h_total(d_arr, generic_curve)
        diffs = np.diff(h)
        assert np.all(diffs >= -1e-6), "H_total must increase monotonically with D"


# ---------------------------------------------------------------------------
# Fixture validation: 5% median relative error on VALID mask (up to one scale)
# ---------------------------------------------------------------------------

class TestFixtureValidation:
    def test_h_total_recovery_error(self, all_fixtures):
        """
        Ĥ_total from the densitometry pipeline must match GT h_total within 5%
        median relative error on the VALID mask pixels, up to one global scale.
        """
        errors_per_case: list[float] = []

        for case in all_fixtures:
            scan = case["scan"]        # (H, W, 3) float32 sRGB
            h_gt = case["h_total"]     # (H, W) float32 ground-truth linear exposure
            stock = case["stock"]

            curve = get_film_curve(stock)
            d_physical, _ = scan_to_density(scan, stock=stock)
            h_rec, conf_mask = density_to_h_total(d_physical, curve)

            # Restrict to VALID pixels only
            valid = conf_mask == VALID
            if valid.sum() < 10:
                continue  # Skip degenerate cases (K=1 with nearly all toe/shoulder)

            h_gt_v = h_gt[valid].astype(np.float64)
            h_rec_v = h_rec[valid].astype(np.float64)

            # Optimal scale (least-squares): g = (h_rec · h_gt) / ||h_rec||²
            denom = float(np.dot(h_rec_v, h_rec_v))
            if denom < 1e-12:
                continue
            g = float(np.dot(h_rec_v, h_gt_v)) / denom
            h_scaled = g * h_rec_v

            rel_err = np.abs(h_scaled - h_gt_v) / (h_gt_v + 1e-8)
            median_err = float(np.median(rel_err))
            errors_per_case.append(median_err)

        assert errors_per_case, "No cases had enough VALID pixels to evaluate"

        overall_median = float(np.median(errors_per_case))
        assert overall_median < 0.05, (
            f"Median relative Ĥ_total error {overall_median:.4f} ≥ 0.05 "
            f"(individual case medians: {[f'{e:.4f}' for e in errors_per_case]})"
        )


class TestAppPathDensitometry:
    """Regression: densitometry must survive the app preprocessing path.

    preprocess_negative once fed densitometry a contrast-stretched (and
    possibly display-inverted) rgb, destroying density information (~90%
    H_total error) while the isolated densitometry tests kept passing.
    The pipeline must run on the pristine scan.
    """

    def test_h_total_accuracy_via_preprocess_negative(self):
        import numpy as np
        from PIL import Image
        from synth.generate import generate_case
        from densitometry import VALID
        from app.preprocessing import preprocess_negative

        errs = []
        for seed in (3, 11, 42):
            case = generate_case(seed=seed, ratio=2.5, size=128)
            pil = Image.fromarray((case["scan"] * 255).astype(np.uint8))
            pre = preprocess_negative(pil, stock=case["stock"])
            assert pre.h_total is not None, "densitometry fields not populated"
            m = pre.confidence_mask == VALID
            assert m.any()
            a, b = pre.h_total[m], case["h_total"][m]
            g = float(np.sum(a * b) / (np.sum(a * a) + 1e-12))
            errs.append(float(np.median(np.abs(g * a - b) / (np.abs(b) + 1e-9))))
        # 8-bit PIL round-trip adds ~1-2% on top of the direct-path ~2%
        assert max(errs) < 0.10, f"app-path H_total errors too high: {errs}"


class TestColorDensitometryRGB:
    """WP-8: rgb densitometry functions (additive; loop scalar)."""

    def test_scan_to_density_rgb_shape_and_nonneg(self):
        rng = np.random.default_rng(0)
        scan = rng.uniform(0.1, 0.9, (32, 32, 3)).astype(np.float32)
        d_rgb = scan_to_density_rgb(scan, "Portra 400 (C-41 color)")
        assert d_rgb.shape == (32, 32, 3)
        assert (d_rgb >= 0).all()

    def test_density_to_h_total_rgb_loops_scalar(self, generic_curve):
        # Use scalar curve for all channels to check loop
        from film_physics import ColorNegativeCurves
        d_rgb = np.full((8, 8, 3), 0.5, dtype=np.float32)
        curves = ColorNegativeCurves(generic_curve, generic_curve, generic_curve, (0.,0.,0.))
        h_rgb, conf_rgb = density_to_h_total_rgb(d_rgb, curves)
        assert h_rgb.shape == (8, 8, 3)
        assert conf_rgb.shape == (8, 8, 3)
        # All channels should give identical since same curve
        assert np.allclose(h_rgb[..., 0], h_rgb[..., 1])
        assert np.allclose(h_rgb[..., 1], h_rgb[..., 2])

    def test_combine_confidence_rgb_truth_table(self):
        mask = np.zeros((2, 2, 3), dtype=np.uint8)
        mask[0, 0] = [0, 1, 1]  # has TOE
        mask[0, 1] = [1, 1, 1]
        mask[1, 0] = [1, 1, 2]  # has SHOULDER
        mask[1, 1] = [1, 1, 1]
        out = combine_confidence_rgb(mask)
        assert out[0, 0] == TOE
        assert out[1, 0] == SHOULDER
        assert out[1, 1] == VALID

    def test_color_fixtures_roundtrip_and_mask_offset(self):
        """Rebuilt per WP-8.1 Fix 2: per-ch median rel Ĥ err (one global scale per ch) on VALID,
        bars r,g<=0.05 b<=0.12, VALID frac >=0.5 per ch, on both committed fixtures.
        """
        from pathlib import Path
        from synth.generate import load_case, generate_case
        from densitometry import scan_to_density_rgb, density_to_h_total_rgb, VALID
        from film_physics import get_color_curves
        fixtures_dir = Path(__file__).parent.parent / "synth" / "fixtures"
        for seed in [0, 1]:
            # load scan from committed fixture
            case_fx = load_case(fixtures_dir / f"color_case_{seed:03d}.npz")
            scan = case_fx["scan"]
            stock = "Portra 400 (C-41 color)"
            d_rgb = scan_to_density_rgb(scan, stock)
            curves = get_color_curves(stock)
            h_from, conf = density_to_h_total_rgb(d_rgb, curves)

            # independent ref h_gt: re-generate case with same seed (deterministic)
            case_ref = generate_case(seed=seed, stock=stock, ratio=2.0, size=64, halation=False, jpeg=False)
            h_gt = case_ref["h_total_rgb"]

            # per ch, find VALID, fit one global scale, compute median rel err
            for c, ch_name in enumerate(["r", "g", "b"]):
                m = (conf[..., c] == VALID)
                if not m.any():
                    continue
                a = h_from[..., c][m]
                b = h_gt[..., c][m]
                g = float(np.sum(a * b) / (np.sum(a * a) + 1e-12))  # one global scale
                err = float(np.median(np.abs(g * a - b) / (np.abs(b) + 1e-9)))
                vfrac = float(m.mean())
                if ch_name == "r" or ch_name == "g":
                    assert err <= 0.05, f"{ch_name} err {err} > 0.05 on seed {seed}"
                else:
                    assert err <= 0.12, f"{ch_name} err {err} > 0.12 on seed {seed}"
                assert vfrac >= 0.5, f"{ch_name} VALID frac {vfrac} < 0.5 on seed {seed}"

    def test_gray_input_equivalence(self):
        """Gray-input equivalence (spec test 2, real): mid-gray (0.5); two INDEPENDENT computations.
        color path green (via curves=) vs scalar B&W path; atol=1e-5.
        """
        from film_physics import FilmCurveParams, PiecewiseFilmCurve, ColorNegativeCurves
        p = FilmCurveParams()  # identical for r g b
        curve = PiecewiseFilmCurve(p)
        curves = ColorNegativeCurves(curve, curve, curve, (0., 0., 0.))
        gray = np.full((16, 16, 3), 0.5, dtype=np.float32)
        # independent color path using curves kwarg (bypasses stock lookup)
        d_green_from_color = scan_to_density_rgb(gray, "Portra 400 (C-41 color)", curves=curves)[..., 1]
        # independent scalar B&W path
        d_scalar = scan_to_density(gray, stock="Generic")[0]
        assert np.allclose(d_green_from_color, d_scalar, atol=1e-5)

    def test_app_level_color(self):
        """App-level color test (spec test 5): goes through process_negative (UI entry) itself."""
        from PIL import Image
        from synth.generate import generate_case
        from app.main import process_negative
        case = generate_case(seed=0, stock="Portra 400 (C-41 color)", size=64)
        pil = Image.fromarray((case["scan"] * 255).astype(np.uint8))
        # call UI entry point with color stock
        ret = process_negative(pil, film_stock="Portra 400 (C-41 color)", physics_weight=1.0, perceptual_weight=0.5, num_candidates=3)
        best_state = ret[5]  # from the return tuple
        assert best_state.get("is_color") is True
        # also the preprocessed path inside would have set it


class TestSlopeMask:
    """Reliability-push mask fix: slope-based VALID band (app opt-in)."""

    def test_default_unchanged_and_slope_narrows(self):
        import numpy as np
        from film_physics import get_film_curve
        from densitometry import density_to_h_total, _slope_valid_bounds, VALID

        curve = get_film_curve("Portra 400")
        d = np.linspace(float(curve.d_min), float(curve.d_max), 512).astype(np.float32).reshape(16, 32)
        h_def, m_def = density_to_h_total(d, curve)
        h_exp, m_exp = density_to_h_total(d, curve, mask_mode="density_margin")
        assert np.array_equal(m_def, m_exp) and np.allclose(h_def, h_exp)

        _h, m_slope = density_to_h_total(d, curve, mask_mode="slope")
        # slope-VALID is a strict subset of margin-VALID on this sweep
        assert np.all((m_slope == VALID) <= (m_def == VALID))
        assert (m_slope == VALID).sum() < (m_def == VALID).sum()

        # Independent bound check: the H-span of the slope band is far narrower.
        d_lo, d_hi = _slope_valid_bounds(curve)
        h_lo = 10 ** float(curve.inverse(np.array([d_lo], np.float32))[0])
        h_hi = 10 ** float(curve.inverse(np.array([d_hi], np.float32))[0])
        margin = 0.05 * (float(curve.d_max) - float(curve.d_min))
        hm_lo = 10 ** float(curve.inverse(np.array([float(curve.d_min) + margin], np.float32))[0])
        hm_hi = 10 ** float(curve.inverse(np.array([float(curve.d_max) - margin], np.float32))[0])
        stops_slope = np.log2(h_hi / h_lo)
        stops_margin = np.log2(hm_hi / hm_lo)
        assert stops_slope < 8.0 < stops_margin  # measured: 6.0 vs 11.1 on Portra