baselines/double_dip.py

"""Double-DIP baseline (WP-6).

Two small untrained hourglass CNNs (deep image priors) map fixed noise to the
two layers. Optimized per-image against the film forward model using
HybridFilmLoss(..., perceptual_weight=0) so the objective is exactly the
masked density residual + exclusion + balance + naturalness (DIP arch provides
the naturalness prior; LPIPS is deliberately not used in the inner loop).

Binding per MASTERPLAN WP-6 execution spec.
"""

from __future__ import annotations

import argparse
import time
from dataclasses import dataclass, field
from typing import Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image

from app.api_client import SeparationResult
from app.preprocessing import preprocess_negative
from film_physics import PiecewiseFilmCurve, get_film_curve
from hybrid_loss import HybridFilmLoss
from scoring_policy import ScoringPolicy, DEFAULT_POLICY, loss_kwargs

# Reuse (do not reimplement) the established downscale/pad/device helpers from latent opt.
from latent_optimizer import (
    downscale_to_max_side as _downscale,
    pad_to_multiple as _pad_to_multiple,
    pick_device as _pick_device,
    rgb_numpy_to_tensor as _rgb_numpy_to_tensor,
    tensor_to_rgb_numpy as _tensor_to_rgb_numpy,
)


@dataclass
class DoubleDIPConfig:
    """Config for Double-DIP separation (all defaults per binding spec)."""

    iterations: int = 2000
    lr: float = 0.01
    max_side: int = 256
    reg_noise_std: float = 1.0 / 30.0
    seed: int = 0
    # WP-13.1: scoring policy (DEFAULT = byte-identical bench path; app passes APP_POLICY)
    policy: ScoringPolicy = field(default_factory=lambda: DEFAULT_POLICY)
    # WP-13.1 F: L2 warm-start iterations before physics phase (opt-in via warm_start arg)
    warm_iters: int = 300


def _build_dip_net() -> nn.Module:
    """Small hourglass CNN (3 levels, channels 16/32/64, refl pad, LeakyReLU, bilinear up, sigmoid).

    Input: (B, 32, H, W) noise in ~[0, 0.1]. Output: (B, 3, H, W) RGB in [0, 1].
    Additive skips at matching decoder stages for a standard DIP hourglass shape.
    """

    class Hourglass(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            # Encoder stages (input noise depth = 32)
            self.e0 = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(32, 16, 3),
                nn.LeakyReLU(0.2, inplace=True),
            )
            self.e1 = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(16, 16, 3),
                nn.LeakyReLU(0.2, inplace=True),
            )
            self.down1 = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(16, 32, 3, stride=2),
                nn.LeakyReLU(0.2, inplace=True),
            )
            self.e2 = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(32, 32, 3),
                nn.LeakyReLU(0.2, inplace=True),
            )
            self.down2 = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(32, 64, 3, stride=2),
                nn.LeakyReLU(0.2, inplace=True),
            )
            self.e3 = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(64, 64, 3),
                nn.LeakyReLU(0.2, inplace=True),
            )
            # Decoder: bilinear up + conv; additive skips from encoder features
            self.up2 = lambda x: F.interpolate(
                x, scale_factor=2, mode="bilinear", align_corners=False
            )
            self.d2 = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(64, 32, 3),
                nn.LeakyReLU(0.2, inplace=True),
            )
            self.up1 = lambda x: F.interpolate(
                x, scale_factor=2, mode="bilinear", align_corners=False
            )
            self.d1 = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(32, 16, 3),
                nn.LeakyReLU(0.2, inplace=True),
            )
            self.d0 = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(16, 3, 3),
                nn.Sigmoid(),
            )

        def forward(self, z: torch.Tensor) -> torch.Tensor:
            x0 = self.e0(z)
            x1 = self.e1(x0)
            x = self.down1(x1)
            x2 = self.e2(x)
            x = self.down2(x2)
            x3 = self.e3(x)
            x = self.up2(x3)
            x = self.d2(x) + F.interpolate(
                x2, size=x.shape[-2:], mode="bilinear", align_corners=False
            )
            x = self.up1(x)
            x = self.d1(x) + F.interpolate(
                x1, size=x.shape[-2:], mode="bilinear", align_corners=False
            )
            x = self.d0(x)
            return x

    return Hourglass()


def double_dip_separate(
    rgb: np.ndarray,
    log_exposure: torch.Tensor,
    density: Optional[np.ndarray],
    confidence_mask: Optional[np.ndarray],
    film_curve: PiecewiseFilmCurve,
    config: Optional[DoubleDIPConfig] = None,
    warm_start: Optional[Tuple[np.ndarray, np.ndarray]] = None,
) -> Optional[SeparationResult]:
    """Run Double-DIP separation.

    Consumes PreprocessedNegative fields (rgb, log_exposure, density, confidence_mask).
    If density is None the source is skipped (return None) — never fall back to legacy.
    Uses HybridFilmLoss(physics=1, perceptual=0) + DIP architecture as naturalness prior.
    Returns one SeparationResult at original resolution (float32 RGB [0,1]).

    WP-13.1 F: optional ``warm_start=(A, B)`` runs ``config.warm_iters`` pure-L2
    reconstruction steps first, then the standard physics objective. Best-snapshot
    is taken only over the physics phase.
    """
    if density is None or confidence_mask is None:
        return None
    if config is None:
        config = DoubleDIPConfig()

    torch.manual_seed(config.seed)
    device = _pick_device()

    orig_h, orig_w = rgb.shape[:2]

    # Downscale observed targets (same pattern as latent_optimizer.refine)
    obs_rgb_t = _downscale(_rgb_numpy_to_tensor(rgb).to(device), config.max_side)
    obs_log_t = _downscale(log_exposure.float().to(device), config.max_side)

    pre_h, pre_w = obs_rgb_t.shape[-2:]

    # Density / mask downscale + pad exactly as refine does
    dens_t = (
        torch.from_numpy(density.astype(np.float32))
        .unsqueeze(0)
        .unsqueeze(0)
        .to(device)
    )
    dens_t = F.interpolate(dens_t, size=(pre_h, pre_w), mode="bilinear", antialias=True)

    conf_t = (
        torch.from_numpy(confidence_mask.astype(np.float32))
        .unsqueeze(0)
        .unsqueeze(0)
        .to(device)
    )
    conf_t = F.interpolate(conf_t, size=(pre_h, pre_w), mode="nearest").round().to(torch.long)

    # Pad working tensors
    obs_rgb_t, unpadded = _pad_to_multiple(obs_rgb_t)
    obs_log_t, _ = _pad_to_multiple(obs_log_t)
    dens_t, _ = _pad_to_multiple(dens_t)
    conf_t, _ = _pad_to_multiple(conf_t, value=0, mode="constant")

    h_work, w_work = obs_rgb_t.shape[-2:]

    # Optional warm-start targets (downscale + pad to work size)
    warm_a_t: Optional[torch.Tensor] = None
    warm_b_t: Optional[torch.Tensor] = None
    warm_iters = 0
    if warm_start is not None:
        wa_np, wb_np = warm_start
        wa_t = _downscale(_rgb_numpy_to_tensor(wa_np.astype(np.float32)).to(device), config.max_side)
        wb_t = _downscale(_rgb_numpy_to_tensor(wb_np.astype(np.float32)).to(device), config.max_side)
        # Match spatial size of obs work grid before pad
        if wa_t.shape[-2:] != (pre_h, pre_w):
            wa_t = F.interpolate(wa_t, size=(pre_h, pre_w), mode="bilinear", align_corners=False)
            wb_t = F.interpolate(wb_t, size=(pre_h, pre_w), mode="bilinear", align_corners=False)
        warm_a_t, _ = _pad_to_multiple(wa_t)
        warm_b_t, _ = _pad_to_multiple(wb_t)
        warm_iters = int(config.warm_iters)

    # Two independent nets + fixed noise inputs
    net_a = _build_dip_net().to(device)
    net_b = _build_dip_net().to(device)

    z_a = (torch.rand(1, 32, h_work, w_work, device=device) * 0.1).detach()
    z_b = (torch.rand(1, 32, h_work, w_work, device=device) * 0.1).detach()

    # Objective: reuse forward_tensor with perceptual=0 (LPIPS skipped; exact WP-3 terms).
    # Deep-copy the curve: .to(device) moves nn.Module state in place, and the caller's
    # curve is shared with CPU-side ranking (mutating it would crash rank_candidates).
    import copy

    # Policy-driven objective (WP-13.1): same loss_kwargs as ranking when APP_POLICY.
    loss_fn = HybridFilmLoss(
        film_curve=copy.deepcopy(film_curve),
        physics_weight=1.0,
        perceptual_weight=0.0,
        **loss_kwargs(config.policy),
    ).to(device)

    # Snapshot best (standard DIP: return lowest-loss iterate, not final)
    with torch.no_grad():
        ia = net_a(z_a)
        ib = net_b(z_b)
        init_loss = float(
            loss_fn.forward_tensor(
                obs_log_t, obs_rgb_t, ia, ib, density=dens_t, confidence_mask=conf_t
            ).item()
        )

    optimizer = torch.optim.Adam(
        list(net_a.parameters()) + list(net_b.parameters()), lr=config.lr
    )

    # Best-snapshot only over the physics phase (WP-13.1 F).
    if warm_iters > 0:
        best_loss = float("inf")
        best_a = ia.detach().clone()
        best_b = ib.detach().clone()
    else:
        best_loss = init_loss
        best_a = ia.detach().clone()
        best_b = ib.detach().clone()

    reg = float(config.reg_noise_std)
    total_iters = int(config.iterations)
    t0 = time.time()
    for it in range(total_iters):
        optimizer.zero_grad()
        za = z_a + reg * torch.randn_like(z_a)
        zb = z_b + reg * torch.randn_like(z_b)
        ia = net_a(za)
        ib = net_b(zb)
        if warm_a_t is not None and it < warm_iters:
            # Pure L2 reconstruction of warm pair (no physics)
            loss = ((ia - warm_a_t) ** 2).mean() + ((ib - warm_b_t) ** 2).mean()
            loss.backward()
            optimizer.step()
            continue
        loss = loss_fn.forward_tensor(
            obs_log_t, obs_rgb_t, ia, ib, density=dens_t, confidence_mask=conf_t
        )
        loss.backward()
        optimizer.step()

        lv = float(loss.item())
        if lv < best_loss:
            best_loss = lv
            best_a = ia.detach().clone()
            best_b = ib.detach().clone()

    dt = time.time() - t0

    # Unpad then bilinear-upscale outputs to original input resolution
    with torch.no_grad():
        ha, wa = unpadded
        ba = best_a[..., :ha, :wa]
        bb = best_b[..., :ha, :wa]
        if (ha, wa) != (orig_h, orig_w):
            ba = F.interpolate(ba, size=(orig_h, orig_w), mode="bilinear", align_corners=False)
            bb = F.interpolate(bb, size=(orig_h, orig_w), mode="bilinear", align_corners=False)
        a_np = _tensor_to_rgb_numpy(ba)
        b_np = _tensor_to_rgb_numpy(bb)

    warm_note = f", warm_iters={warm_iters}" if warm_iters > 0 else ""
    msg = (
        f"Double-DIP baseline ({config.iterations} iters{warm_note}, {dt:.1f}s, "
        f"best_loss={best_loss:.6f}, init_loss={init_loss:.6f}, device={device.type})"
    )
    cid = f"dip_i{config.iterations}"
    if warm_iters > 0:
        cid = f"dip_i{config.iterations}_w{warm_iters}"
    return SeparationResult(
        image_a=a_np.astype(np.float32),
        image_b=b_np.astype(np.float32),
        method="deep_prior",
        message=msg,
        candidate_id=cid,
        diagnostics={
            "init_loss": init_loss,
            "best_loss": best_loss,
            "runtime_s": dt,
            "device": device.type,
            "iterations": int(config.iterations),
            "warm_iters": warm_iters,
        },
    )


def run_bench(iterations: int = 2000, fixtures_dir: str | None = None, limit: int = 6) -> None:
    """Real bench used for WP-6 accept: DIP vs demo heuristics on the K=2 fixtures (up to `limit`).

    Invoked via: python -m baselines.double_dip --bench [--iters N] [--fixtures-dir DIR]
    Output table is pasted verbatim into MASTERPLAN WP-6 Result note; the note must
    state the iteration count so the table is reproducible from committed code.
    """
    from synth.evaluation import score_pair
    from synth.generate import load_fixtures, select_k2_cases

    from app.api_client import generate_demo_candidates

    dir_note = f" --fixtures-dir {fixtures_dir}" if fixtures_dir else ""
    print(f"Double-DIP --bench ({iterations} iters{dir_note}, limit {limit}): DIP vs demo heuristics on the K=2 fixtures")
    cases = load_fixtures(fixtures_dir)
    k2_cases = select_k2_cases(cases, limit)
    print(f"Using {len(k2_cases)} K=2 fixtures from {fixtures_dir or 'synth/fixtures'}")

    rows: list[tuple] = []
    for idx, case in enumerate(k2_cases):
        scan = case["scan"]
        pil = Image.fromarray((scan * 255.0).clip(0, 255).astype(np.uint8))
        pre = preprocess_negative(pil)
        if pre.density is None or pre.confidence_mask is None:
            print(f"  case {idx} (seed {case['seed']}): skipped — no density/confidence")
            continue
        stock = case.get("stock", "Generic")
        curve = get_film_curve(stock)

        t0 = time.time()
        cfg = DoubleDIPConfig(iterations=iterations, seed=42 + idx)
        res = double_dip_separate(
            pre.rgb, pre.log_exposure, pre.density, pre.confidence_mask, curve, cfg
        )
        dt = time.time() - t0
        if res is None:
            continue
        dip_sc = score_pair(case["gt_a"], case["gt_b"], res.image_a, res.image_b)
        dip_lp = float(dip_sc.get("lpips", float("nan")))

        # Demo heuristics (same positive_rgb fed to generate_demo_candidates)
        demos = generate_demo_candidates(pre.rgb, num_candidates=5)
        heu_lp = float("inf")
        for d in demos:
            sc = score_pair(case["gt_a"], case["gt_b"], d.image_a, d.image_b)
            l = float(sc.get("lpips", float("inf")))
            if l < heu_lp:
                heu_lp = l

        rows.append((idx, case["seed"], dip_lp, heu_lp, dt))
        print(f"case{idx} seed{case['seed']}: DIP={dip_lp:.4f} heu_best={heu_lp:.4f} t={dt:.1f}s")

    print("\nPer-case table (DIP vs best-of-heuristics demo):")
    print("| case | seed | DIP LPIPS | best-heu LPIPS | runtime_s |")
    print("|------|------|-----------|----------------|-----------|")
    for r in rows:
        print(f"| {r[0]} | {r[1]} | {r[2]:.4f} | {r[3]:.4f} | {r[4]:.1f} |")
    if rows:
        md = np.mean([r[2] for r in rows])
        mh = np.mean([r[3] for r in rows])
        print(f"\nMean per-layer LPIPS: DIP={md:.4f}  best-heu={mh:.4f}")
    print("Bench complete.")


def main() -> None:
    parser = argparse.ArgumentParser(description="Double-DIP baseline (WP-6)")
    parser.add_argument(
        "--bench", action="store_true", help="Run on the K=2 fixtures (up to --limit) and compare LPIPS to demo heuristics"
    )
    parser.add_argument(
        "--iters", type=int, default=2000, help="DIP iterations per case for the bench (default: spec 2000)"
    )
    parser.add_argument(
        "--fixtures-dir", type=str, default=None, dest="fixtures_dir",
        help="Directory with case_*.npz fixtures (default: synth/fixtures; for 256px re-measure)"
    )
    parser.add_argument(
        "--limit", type=int, default=6, help="Max K=2 cases to process (default 6; for 50-case bench)"
    )
    args = parser.parse_args()
    if args.bench:
        run_bench(iterations=args.iters, fixtures_dir=args.fixtures_dir, limit=args.limit)
    else:
        print("Double-DIP baseline ready. Use --bench for the fixture comparison (slow).")


if __name__ == "__main__":
    main()