tests/test_double_dip.py

"""Fast offline tests for WP-6 Double-DIP baseline.

- Smoke: tiny iters + 32x32 crop → valid SeparationResult + final loss < initial (teeth-proven).
- Registration: include_deep_prior=True yields method=="deep_prior" candidate that ranks.
- density=None path is skipped (no candidate added).

Full 2000-iter bench is CLI-only and marked slow (not run in CI).
"""

from __future__ import annotations

import numpy as np
import pytest
from PIL import Image

from app.api_client import SeparationResult
from app.preprocessing import preprocess_negative
from app.scoring import rank_candidates
from film_physics import get_film_curve

# Direct import for core tests (fast path)
from baselines.double_dip import DoubleDIPConfig, double_dip_separate


def _load_fixture_as_pil(idx: int = 0) -> Image.Image:
    from pathlib import Path

    fix = sorted(Path("synth/fixtures").glob("case_*.npz"))[idx]
    d = np.load(fix)
    arr = (d["scan"] * 255).clip(0, 255).astype(np.uint8)
    return Image.fromarray(arr)


def test_double_dip_smoke():
    """Tiny config returns valid shapes/dtypes/ranges and improves loss (teeth: fails if opt disabled)."""
    # Use full small fixture (64x64); DIP downs internally via max_side. Crop was causing densitometry soft-fail on tiny input.
    pil = _load_fixture_as_pil(0)
    pre = preprocess_negative(pil)
    curve = get_film_curve("Generic")
    cfg = DoubleDIPConfig(iterations=30, max_side=64, seed=0)

    res = double_dip_separate(
        pre.rgb, pre.log_exposure, pre.density, pre.confidence_mask, curve, cfg
    )
    assert res is not None
    assert isinstance(res, SeparationResult)
    assert res.method == "deep_prior"
    assert res.candidate_id.startswith("dip_i")
    assert res.image_a.shape == res.image_b.shape == pre.rgb.shape
    assert res.image_a.dtype == res.image_b.dtype == np.float32
    assert 0.0 <= float(res.image_a.min()) <= float(res.image_a.max()) <= 1.0
    assert 0.0 <= float(res.image_b.min()) <= float(res.image_b.max()) <= 1.0

    # Structured diagnostics (not message parsing) carry the loss improvement
    assert res.diagnostics is not None
    init_l = res.diagnostics["init_loss"]
    best_l = res.diagnostics["best_loss"]
    assert best_l < init_l, f"DIP did not improve loss: init={init_l} best={best_l}"


def test_double_dip_registration(monkeypatch):
    """The REAL integration surface: generate_candidates(include_deep_prior=True) registers a
    deep_prior candidate (tiny config monkeypatched in) that rank_candidates then scores.
    """
    import baselines.double_dip as dd
    from app.api_client import generate_candidates

    pil = _load_fixture_as_pil(0)
    pre = preprocess_negative(pil)
    curve = get_film_curve("Generic")

    # Tiny config so the wiring test stays fast; the helper's lazy import picks this up.
    monkeypatch.setattr(
        dd, "DoubleDIPConfig", lambda **kw: DoubleDIPConfig(iterations=12, max_side=64, seed=7)
    )

    cands, mode = generate_candidates(
        pre.rgb,
        num_candidates=2,
        h_total=pre.h_total,
        confidence_mask=pre.confidence_mask,
        density=pre.density,
        log_exposure=pre.log_exposure,
        include_deep_prior=True,
        film_curve=curve,
    )
    assert mode == "demo"
    dips = [c for c in cands if c.method == "deep_prior"]
    assert dips, "generate_candidates(include_deep_prior=True) did not register a deep_prior candidate"

    # The mixed pool participates in normal ranking
    ranked = rank_candidates(
        candidates=cands,
        observed_log_exposure=pre.log_exposure,
        observed_rgb=pre.rgb,
        film_curve=curve,
        physics_weight=1.0,
        perceptual_weight=0.0,
        density=pre.density,
        confidence_mask=pre.confidence_mask,
    )
    assert any(r.separation.method == "deep_prior" for r in ranked)


def test_double_dip_skips_without_density():
    """density=None (or conf=None) produces no deep_prior candidate (same rule as demix)."""
    pil = _load_fixture_as_pil(0)
    pre = preprocess_negative(pil)

    # Direct path (generate wiring with the kw args lands in commit 2)
    res = double_dip_separate(
        pre.rgb, pre.log_exposure, density=None, confidence_mask=pre.confidence_mask, film_curve=get_film_curve("Generic")
    )
    assert res is None