synth/real_protocol.py

"""WP-13 real-photo scoring protocol (privacy-binding).

READS photos from --photos-dir only. Never commits images. Optional saves go
under real_outputs/ (gitignored). Result note records NUMBERS/tables only.

Usage:
  .venv/bin/python -m synth.real_protocol \\
      --photos-dir /path/to/photos \\
      --photos 127,124,131 \\
      [--limit N] [--dip-iters 2000] [--no-save]
"""

from __future__ import annotations

import argparse
import time
from pathlib import Path
from typing import List, Optional, Sequence

import numpy as np
from PIL import Image

from app.api_client import SeparationResult, generate_demo_candidates
from app.demix import residual_demix, stub_cleanup, DemixConfig, analyze_scan
from app.preprocessing import preprocess_negative
from app.scoring import rank_candidates, rank_then_merge_asymmetric
from baselines.double_dip import DoubleDIPConfig, double_dip_separate
from film_physics import get_film_curve
from hybrid_loss import _luminance_from_rgb
from scoring_policy import APP_POLICY


def _match_photos(photos_dir: Path, names: Optional[Sequence[str]], limit: Optional[int]) -> List[Path]:
    files = sorted(
        p for p in photos_dir.iterdir()
        if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp"}
    )
    if names:
        # Match substring (e.g. "127" → EdwardFal...-127.jpg)
        selected: List[Path] = []
        for name in names:
            hits = [p for p in files if name in p.stem or name in p.name]
            if not hits:
                raise FileNotFoundError(f"No photo matching '{name}' in {photos_dir}")
            selected.append(hits[0])
        files = selected
    if limit is not None:
        files = files[: int(limit)]
    return files


def _std_ratio(img_a: np.ndarray, img_b: np.ndarray, obs: np.ndarray) -> tuple[float, float, float]:
    la = _luminance_from_rgb(img_a)
    lb = _luminance_from_rgb(img_b)
    lo = _luminance_from_rgb(obs)
    std_o = float(np.std(lo)) + 1e-12
    ra = float(np.std(la)) / std_o
    rb = float(np.std(lb)) / std_o
    return ra, rb, min(ra, rb)


def _run_one(
    path: Path,
    stock: str,
    dip_iters: int,
    save_dir: Optional[Path],
    dip_warm_start: bool = False,
    asym: bool = False,
) -> str:
    t_all = time.time()
    img = Image.open(path)
    pre = preprocess_negative(
        img,
        stock=stock,
        scan_type="positive",
        scan_calibration="auto_exposed",
        auto_trim=False,
    )
    curve = get_film_curve(stock)
    lines: List[str] = []
    lines.append(f"## Photo `{path.name}`")
    lines.append(
        f"- working size: {pre.rgb.shape[1]}×{pre.rgb.shape[0]} "
        f"(orig {pre.original_size[0]}×{pre.original_size[1]})"
    )
    lines.append(f"- density present: {pre.density is not None}; "
                 f"mask present: {pre.confidence_mask is not None}")
    lines.append("")

    candidates: List[SeparationResult] = []
    timings: dict[str, float] = {}

    # Heuristics (demo)
    t0 = time.time()
    demos = generate_demo_candidates(pre.rgb, num_candidates=3)
    timings["heuristics"] = time.time() - t0
    candidates.extend(demos)

    # Demix stub
    if pre.h_total is not None and pre.confidence_mask is not None:
        t0 = time.time()
        try:
            analysis = analyze_scan(pre.rgb, vlm=None)
            d = residual_demix(
                pre.rgb,
                pre.h_total,
                pre.confidence_mask,
                stub_cleanup,
                analysis,
                DemixConfig(iterations=2, strength=0.55),
                method="demix_stub",
            )
            candidates.append(d)
        except Exception as exc:
            lines.append(f"- demix failed: {exc}")
        timings["demix"] = time.time() - t0

    # Optional warm-start: rank structured pool first, take top (A,B)
    warm_pair = None
    if dip_warm_start and candidates:
        pre_rank = rank_candidates(
            candidates=candidates,
            observed_log_exposure=pre.log_exposure,
            observed_rgb=pre.rgb,
            film_curve=curve,
            physics_weight=1.0,
            perceptual_weight=0.5,
            density=pre.density,
            confidence_mask=pre.confidence_mask,
            policy=APP_POLICY,
        )
        if len(pre_rank) > 0:
            top = pre_rank.ranked[0]
            warm_pair = (top.separation.image_a, top.separation.image_b)
            lines.append(
                f"- warm-start from `{top.candidate_id}` "
                f"(pre-rank total={top.score.total_loss:.4f})"
            )

    # DIP with APP_POLICY (same objective as ranking) — before single rank
    if pre.density is not None and pre.confidence_mask is not None:
        t0 = time.time()
        try:
            cfg = DoubleDIPConfig(
                iterations=int(dip_iters),
                seed=0,
                policy=APP_POLICY,
            )
            dip = double_dip_separate(
                pre.rgb,
                pre.log_exposure,
                pre.density,
                pre.confidence_mask,
                curve,
                config=cfg,
                warm_start=warm_pair,
            )
            if dip is not None:
                candidates.append(dip)
        except Exception as exc:
            lines.append(f"- DIP failed: {exc}")
        timings["dip"] = time.time() - t0
    else:
        lines.append("- DIP skipped (no density/mask)")

    # WP-14.1: rank base ONCE; score-only-merge asym (no second full LPIPS pass)
    t0 = time.time()
    asym_notes: list[str] = []
    ranked, _asym_api = rank_then_merge_asymmetric(
        candidates,
        pre,
        film_curve=curve,
        physics_weight=1.0,
        perceptual_weight=0.5,
        policy=APP_POLICY,
        enable_asymmetric=bool(asym),
        anchor="auto",
        complete=False,  # offline protocol: physics only
        status_notes=asym_notes,
        debug=False,
    )
    for n in asym_notes:
        lines.append(f"- asym: {n}")
    timings["rank"] = time.time() - t0
    timings["asym"] = 0.0  # fold into rank helper; wall is rank total

    lines.append(
        f"- candidates generated: {len(candidates)}; "
        f"ranked: {len(ranked)}; "
        f"flat rejected: {ranked.rejected_count}"
    )
    lines.append(
        f"- wall times (s): heuristics={timings.get('heuristics', 0):.2f}, "
        f"demix={timings.get('demix', 0):.2f}, "
        f"asym={timings.get('asym', 0):.2f}, "
        f"dip={timings.get('dip', 0):.2f}, "
        f"rank={timings.get('rank', 0):.2f}, "
        f"total={time.time() - t_all:.2f}"
    )
    lines.append("")
    lines.append(
        "| rank | id | method | total | phys | grad | excl | bal | perc | "
        "a | r | stdA/obs | stdB/obs | min_ratio |"
    )
    lines.append("|---:|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|")

    for item in ranked:
        s = item.score
        sep = item.separation
        ra, rb, rmin = _std_ratio(sep.image_a, sep.image_b, pre.rgb)
        lines.append(
            f"| {item.rank} | {item.candidate_id} | {sep.method} | "
            f"{s.total_loss:.4f} | {s.physics_loss:.4f} | {s.physics_grad_loss:.4f} | "
            f"{s.exclusivity_loss:.4f} | {s.balance_loss:.4f} | {s.perceptual_loss:.4f} | "
            f"{s.affine_a:.3f} | {s.fitted_r:.3f} | {ra:.3f} | {rb:.3f} | {rmin:.3f} |"
        )

    # Accept diagnostics (numbers only)
    best = ranked[0] if ranked else None
    if best is not None:
        ra, rb, rmin = _std_ratio(best.separation.image_a, best.separation.image_b, pre.rgb)
        lines.append("")
        lines.append(
            f"**Best:** `{best.candidate_id}` method={best.separation.method} "
            f"total={best.score.total_loss:.4f} a={best.score.affine_a:.3f} "
            f"r={best.score.fitted_r:.3f} min_std_ratio={rmin:.3f}"
        )
        # DIP vs best heuristic r
        dip_items = [r for r in ranked if r.separation.method == "deep_prior"]
        heur_items = [r for r in ranked if r.separation.method.startswith("demo")]
        if dip_items and heur_items:
            dip_r = dip_items[0].score.fitted_r
            best_heur_r = max(h.score.fitted_r for h in heur_items)
            dra, drb, drmin = _std_ratio(
                dip_items[0].separation.image_a,
                dip_items[0].separation.image_b,
                pre.rgb,
            )
            lines.append(
                f"**DIP accept check:** std_ratio_min={drmin:.4f} "
                f"(need >0.05 for non-flat); fitted_r={dip_r:.4f} vs best_heuristic_r={best_heur_r:.4f}"
            )

    if save_dir is not None and best is not None:
        save_dir.mkdir(parents=True, exist_ok=True)
        stem = path.stem
        for label, arr in (
            ("A", best.separation.image_a),
            ("B", best.separation.image_b),
        ):
            out = (np.clip(arr, 0, 1) * 255).astype(np.uint8)
            Image.fromarray(out).save(save_dir / f"{stem}_{label}.png")

    lines.append("")
    return "\n".join(lines)


def main(argv: Optional[Sequence[str]] = None) -> int:
    p = argparse.ArgumentParser(description="WP-13 real-photo scoring protocol")
    p.add_argument("--photos-dir", type=str, required=True)
    p.add_argument("--limit", type=int, default=None)
    p.add_argument("--photos", type=str, default=None, help="Comma-separated name substrings")
    p.add_argument("--stock", type=str, default="Portra 400")
    p.add_argument("--dip-iters", type=int, default=2000)
    p.add_argument(
        "--dip-warm-start",
        action="store_true",
        help="Warm-start DIP from top pre-ranked structured candidate (WP-13.1 F)",
    )
    p.add_argument(
        "--asym",
        action="store_true",
        help="WP-14: add offline asymmetric subtraction candidate (asym_sub)",
    )
    p.add_argument("--no-save", action="store_true")
    p.add_argument(
        "--output-dir",
        type=str,
        default="real_outputs",
        help="Gitignored dir for optional image saves (never commit)",
    )
    args = p.parse_args(argv)

    photos_dir = Path(args.photos_dir)
    if not photos_dir.is_dir():
        raise SystemExit(f"photos-dir not found: {photos_dir}")

    names = [s.strip() for s in args.photos.split(",")] if args.photos else None
    paths = _match_photos(photos_dir, names, args.limit)
    if not paths:
        raise SystemExit("No photos matched")

    save_dir = None if args.no_save else Path(args.output_dir)
    print(
        f"# WP-13/14 real protocol — {len(paths)} photo(s), stock={args.stock}, "
        f"dip_iters={args.dip_iters}, warm_start={args.dip_warm_start}, "
        f"asym={args.asym}"
    )
    print(f"# photos-dir={photos_dir} (READ ONLY); saves={'off' if save_dir is None else save_dir}")
    print()

    for path in paths:
        print(
            _run_one(
                path,
                stock=args.stock,
                dip_iters=args.dip_iters,
                save_dir=save_dir,
                dip_warm_start=bool(args.dip_warm_start),
                asym=bool(args.asym),
            )
        )
        print()

    return 0


if __name__ == "__main__":
    raise SystemExit(main())