app/main.py

"""Gradio web interface for double-exposure negative recovery."""

from __future__ import annotations

import os
import traceback

import gradio as gr
import numpy as np
from dotenv import load_dotenv
from PIL import Image

from film_physics import get_film_curve, list_film_stocks, get_color_curves, list_color_stocks, COLOR_STOCK_PRESETS
from app.api_client import generate_candidates, result_to_pil_pair
from app.preprocessing import preprocess_negative, to_pil
from app.scoring import (
    format_ranking_table,
    format_score_summary,
    rank_candidates,
    rank_then_merge_asymmetric,
    score_separation,
)
from scoring_policy import APP_POLICY
from app.scribble import PALETTE_RGBA as _PALETTE_RGBA

# HF Spaces free tier runs on ZeroGPU, which requires at least one
# @spaces.GPU-decorated function at startup and grants CUDA only inside such
# functions. The spaces package is a no-op on any other hardware (local dev,
# CPU Spaces); the fallback keeps environments without the package working.
try:
    import spaces as _hf_spaces
except ImportError:  # pragma: no cover — local envs without the spaces package
    class _hf_spaces:  # type: ignore[no-redef]
        @staticmethod
        def GPU(fn=None, **_kwargs):
            if callable(fn):
                return fn
            return lambda f: f


load_dotenv()

FILM_STOCKS = list(list_film_stocks()) + list(list_color_stocks())


def _get_ranking_curve(stock: str):
    """Resolver for ranking: for color stocks, use the green channel curve (scalar path)."""
    if stock in COLOR_STOCK_PRESETS:
        return get_color_curves(stock).g
    return get_film_curve(stock)


# WP-18 D4a: ONE source of truth for UI-label -> preprocess param mapping,
# used by process_negative AND restore_best_scene's standalone path (the two
# inline copies had already drifted on capitalized keys).
_ST_MAP = {"Auto": "auto", "Positive": "positive", "Negative": "negative"}
_CAL_MAP = {
    "auto-exposed": "auto_exposed",
    "Auto-exposed": "auto_exposed",
    "calibrated-linear": "linear",
    "Calibrated-linear": "linear",
    "auto_exposed": "auto_exposed",
    "linear": "linear",
}


def _map_scan_params(scan_type: str, scan_calibration: str) -> tuple[str, str]:
    st = _ST_MAP.get(scan_type, (scan_type or "auto").lower())
    cal = _CAL_MAP.get(scan_calibration, "auto_exposed")
    return st, cal


TRANSPARENCY_NOTICE = (
    "> **Note:** Results are best-effort AI reconstructions, not perfect recoveries. "
    "Heavy overlap or similar tones between exposures may limit separation quality."
)


ASYM_CONSENT = (
    "AI completion sends your image to Replicate (and scene analysis to Anthropic) "
    "when API keys are configured. Without keys, only the offline physics subtraction runs."
)


@_hf_spaces.GPU(duration=60)
def process_negative(
    upload: Image.Image | None,
    film_stock: str,
    physics_weight: float,
    perceptual_weight: float,
    num_candidates: int,
    include_deep_prior: bool = False,
    full_res_export: bool = False,
    scan_type: str = "Auto",
    scan_calibration: str = "auto-exposed",
    auto_trim: bool = False,
    asymmetric_recovery: bool = False,
    asymmetric_anchor: str = "Auto",
) -> tuple:
    """Preprocess, generate candidates, rank by hybrid loss, return best result.
    If full_res_export, also compute and return full-res A/B via app/fullres
    using the ORIGINAL upload (not the resized preprocessed.rgb).
    WP-14: optional asymmetric_recovery appends asym_sub (+ asym_fill if keys).
    """
    empty = (None, None, None, None, TRANSPARENCY_NOTICE, None, None, None, None, None, None)

    if upload is None:
        return (
            *empty[:4],
            "Upload a scanned negative to begin.\n\n" + TRANSPARENCY_NOTICE,
            None,
            None,
            None,
            None,
            None,
            None,
        )

    try:
        # WP-11 Fix B/C via the shared map (WP-18 D4a)
        st, cal = _map_scan_params(scan_type, scan_calibration)
        preprocessed = preprocess_negative(
            upload,
            stock=film_stock,
            scan_type=st,
            scan_calibration=cal,
            auto_trim=bool(auto_trim),
            mask_mode="slope",
        )
        positive_pil = _overlay_confidence(preprocessed.rgb, preprocessed.confidence_mask) if preprocessed.confidence_mask is not None else to_pil(preprocessed.rgb)

        film_curve = _get_ranking_curve(film_stock)
        candidates, mode = generate_candidates(
            preprocessed.rgb,
            num_candidates=int(num_candidates),
            h_total=preprocessed.h_total,
            confidence_mask=preprocessed.confidence_mask,
            density=preprocessed.density,
            log_exposure=preprocessed.log_exposure,
            include_deep_prior=bool(include_deep_prior),
            film_curve=film_curve,
            dip_policy=APP_POLICY,
        )

        asym_notes: list[str] = []
        from app.asymmetric import normalize_anchor

        # Gradio boundary: normalize once to Literal auto|a|b
        anchor_key = normalize_anchor(str(asymmetric_anchor))
        # WP-14.1: single base rank; score-only merge; api_contacted for consent
        ranked, asym_api_contacted = rank_then_merge_asymmetric(
            candidates,
            preprocessed,
            film_curve=film_curve,
            physics_weight=physics_weight,
            perceptual_weight=perceptual_weight,
            policy=APP_POLICY,
            enable_asymmetric=bool(asymmetric_recovery),
            anchor=anchor_key,
            complete=bool(asymmetric_recovery),  # soft-fails offline without key
            status_notes=asym_notes,
            debug=False,  # never pass True — scalar diagnostics only
        )
        # Keep candidates list in sync for VLM note / gallery identity
        if asymmetric_recovery:
            for item in ranked:
                if item.candidate_id in ("asym_sub", "asym_fill"):
                    if not any(c.candidate_id == item.candidate_id for c in candidates):
                        candidates = list(candidates) + [item.separation]

        best = ranked[0]
        img_a, img_b = result_to_pil_pair(best.separation)
        recombined_pil = to_pil(best.score.recombined_rgb)

        invert_note = (
            " _(auto-inverted from negative scan)_" if preprocessed.was_inverted else ""
        )
        mode_label = {
            "demo": "Demo (no API key)",
            "replicate": "Replicate API",
            "demo_fallback": "Demo fallback (API error)",
        }.get(mode, mode)

        score_text = format_score_summary(
            best, len(ranked), physics_weight, perceptual_weight
        )
        ranking_table = format_ranking_table(ranked)

        dens_note = ""
        if preprocessed.density is None:
            dens_note = "\n\n**Warning:** densitometry unavailable — physics running in legacy mode (circular path)."

        vlm_note = ""
        if any(c.method.startswith("demix") for c in candidates):
            try:
                # WP-5.1 Fix 7: prefer carried scan_analysis to avoid redundant VLM call
                a = next((c.scan_analysis for c in candidates if getattr(c, "scan_analysis", None) is not None), None)
                if a is None:
                    from app.demix import analyze_scan, anthropic_vlm
                    import os
                    vlm = anthropic_vlm if os.environ.get("ANTHROPIC_API_KEY") else None
                    a = analyze_scan(preprocessed.rgb, vlm=vlm)
                vlm_note = (
                    f"\n\n**VLM analysis:** scene A: {a.scene_a_description[:80]}; "
                    f"scene B: {a.scene_b_description[:80]}; k_judgment={a.k_judgment:.2f}"
                )
            except Exception:
                pass

        guard_note = ""
        rejected = ranked.rejected_count
        if rejected > 0:
            guard_note = f"\n\n**{rejected} degenerate candidate(s) discarded** (flat-layer guard)."

        asym_note = ""
        if asymmetric_recovery:
            asym_note = f"\n\n**Asymmetric recovery:** {ASYM_CONSENT}"
            if asym_notes:
                asym_note += "\n\n" + "\n\n".join(f"- {n}" for n in asym_notes)
            # WP-14.1 P1: consent from api_contacted flag — never note-string matching
            if asym_api_contacted:
                asym_note += f"\n\n*{ASYM_CONSENT}*"

        status = (
            f"{TRANSPARENCY_NOTICE}\n\n"
            f"**Mode:** {mode_label}\n\n"
            f"**Film stock:** {film_stock}{invert_note}\n\n"
            f"**Original size:** {preprocessed.original_size[0]}×"
            f"{preprocessed.original_size[1]} px\n\n"
            f"{score_text}\n\n"
            f"### All candidates\n{ranking_table}{dens_note}{vlm_note}{guard_note}{asym_note}"
        )

        best_state = _build_best_state(
            preprocessed, best.separation, film_stock, physics_weight, perceptual_weight
        )

        # WP-4: prepare gallery data (list of (thumbnail, caption)) and confidence viz
        gallery_data: list[tuple[Image.Image, str]] = []
        for item in ranked:
            thumb = _make_candidate_thumbnail(item.separation.image_a, item.separation.image_b)
            cap = f"#{item.rank} {item.candidate_id} (loss={item.score.total_loss:.4f}, {item.separation.method})"
            gallery_data.append((thumb, cap))
        conf_pil = _confidence_mask_to_pil(preprocessed.confidence_mask)

        # WP-12: full-res using ORIGINAL upload (not preprocessed.rgb)
        full_a_pil = None
        full_b_pil = None
        if full_res_export:
            try:
                from app.fullres import upscale_separation, FullResConfig
                from PIL import ImageOps

                # WP-11 post-review: the working images come from the
                # EXIF-transposed (and possibly auto-trimmed) pipeline. The
                # full-res guide must match that geometry or the split-ratio
                # field lands on rotated / border-padded content.
                orig_src = ImageOps.exif_transpose(upload) or upload
                orig_arr = np.asarray(orig_src.convert("RGB"), dtype=np.float32) / 255.0
                if preprocessed.trim_bbox_frac is not None:
                    tf, bf, lf, rf = preprocessed.trim_bbox_frac
                    oh, ow = orig_arr.shape[:2]
                    orig_arr = orig_arr[
                        int(round(tf * oh)) : int(round(bf * oh)),
                        int(round(lf * ow)) : int(round(rf * ow)),
                    ]
                sep = best.separation
                a_f, b_f = upscale_separation(
                    orig_arr,
                    sep.image_a,
                    sep.image_b,
                    stock=film_stock,
                    config=FullResConfig(),
                    positive_source=preprocessed.physics_is_positive,
                    d_min_override=preprocessed.d_min_override_used,
                )
                full_a_pil = to_pil(a_f)
                full_b_pil = to_pil(b_f)
            except Exception as e:
                status += f"\n\n**Full-res export failed:** {e}"

        return positive_pil, recombined_pil, img_a, img_b, status, best_state, ranked, gallery_data, conf_pil, full_a_pil, full_b_pil

    except Exception as exc:
        tb = traceback.format_exc()
        return (
            None,
            None,
            None,
            None,
            f"**Error during processing:** {exc}\n\n```\n{tb}\n```\n\n{TRANSPARENCY_NOTICE}",
            None,
            None,
            None,
            None,
            None,
            None,
        )


@_hf_spaces.GPU(duration=60)
def enhance_best_result(
    best_state: dict | None,
    opt_steps: float,
    opt_lr: float,
) -> tuple:
    """Run gradient-based latent refinement on the current best separation."""
    if not best_state:
        return None, None, "Run **Recover exposures** first, then enhance the best result."

    try:
        from hybrid_loss import HybridFilmLoss
        from latent_optimizer import LatentSpaceOptimizer

        film_curve = _get_ranking_curve(best_state["film_stock"])
        loss_fn = HybridFilmLoss(
            film_curve=film_curve,
            physics_weight=best_state["physics_weight"],
            perceptual_weight=best_state["perceptual_weight"],
        )
        optimizer = LatentSpaceOptimizer(
            hybrid_loss=loss_fn,
            steps=int(opt_steps),
            lr=float(opt_lr),
        )
        result = optimizer.refine(
            recon_a=best_state["image_a"],
            recon_b=best_state["image_b"],
            observed_log_exposure=best_state["log_exposure"],
            observed_rgb=best_state["observed_rgb"],
            observed_density=best_state.get("density"),
            confidence_mask=best_state.get("confidence_mask"),
        )

        refined_a = to_pil(result.refined_a)
        refined_b = to_pil(result.refined_b)

        space = "VAE latent space" if result.used_vae else "pixel space (VAE unavailable)"
        dens_used = getattr(result, "used_density", False)
        dens_note = " (density path)" if dens_used else " (legacy path)"

        # Fix 3: re-score original and refined via the *ranking* objective for headline.
        # WP-13.1 E: scoring uses APP_POLICY; the refine OPTIMIZER objective is still
        # legacy (follow-up WP). Accept-gate makes that safe (can fail to improve but
        # cannot ship worse than the original images).
        film_curve = _get_ranking_curve(best_state["film_stock"])
        scored_orig = score_separation(
            observed_log_exposure=best_state["log_exposure"],
            observed_rgb=best_state["observed_rgb"],
            image_a_rgb=best_state["image_a"],
            image_b_rgb=best_state["image_b"],
            film_curve=film_curve,
            physics_weight=best_state["physics_weight"],
            perceptual_weight=best_state["perceptual_weight"],
            density=best_state.get("density"),
            confidence_mask=best_state.get("confidence_mask"),
            policy=APP_POLICY,
        )
        scored_refined = score_separation(
            observed_log_exposure=best_state["log_exposure"],
            observed_rgb=best_state["observed_rgb"],
            image_a_rgb=result.refined_a,
            image_b_rgb=result.refined_b,
            film_curve=film_curve,
            physics_weight=best_state["physics_weight"],
            perceptual_weight=best_state["perceptual_weight"],
            density=best_state.get("density"),
            confidence_mask=best_state.get("confidence_mask"),
            policy=APP_POLICY,
        )

        if scored_refined.total_loss < scored_orig.total_loss:
            verdict = f"improved (headline score {scored_refined.total_loss:.6f} vs {scored_orig.total_loss:.6f})"
            out_a, out_b = refined_a, refined_b
        else:
            verdict = "did not improve on ranking objective (reverted to original images)"
            out_a = to_pil(best_state["image_a"])
            out_b = to_pil(best_state["image_b"])
        summary = (
            f"**Physics optimization complete** — {result.steps_run} steps in {space}{dens_note}.\n\n"
            f"- Headline (re-scored via ranking objective): before {scored_orig.total_loss:.6f} → after {scored_refined.total_loss:.6f}\n"
            f"- Result: {verdict}\n\n"
            f"- Internal forward_tensor before/after (for reference): {result.initial_loss:.6f} → {result.final_loss:.6f}\n\n"
            "_Headline numbers come from score_separation (the ranking loss); internal losses may differ due to path divergence._"
        )
        return out_a, out_b, summary

    except Exception as exc:
        tb = traceback.format_exc()
        return None, None, f"**Enhancement failed:** {exc}\n\n```\n{tb}\n```"


def promote_candidate(
    evt: gr.SelectData,
    ranked_list: list | None,
    current_status: str,
    current_best_state: dict | None,
) -> tuple:
    """WP-4: when user clicks gallery item, promote that candidate as the new 'best'."""
    if not ranked_list or evt.index is None:
        return gr.update(), gr.update(), gr.update(), gr.update(), gr.update()

    idx = evt.index
    if idx < 0 or idx >= len(ranked_list):
        return gr.update(), gr.update(), gr.update(), gr.update(), gr.update()

    item = ranked_list[idx]  # RankedCandidate
    pil_a, pil_b = result_to_pil_pair(item.separation)
    recomb_pil = to_pil(item.score.recombined_rgb)

    # Sync best_state with promoted candidate for enhance etc.
    # WP-4.1 Fix 4: use shared _build_best_state to prevent schema drift
    if current_best_state:
        class _Pre: pass
        pre = _Pre()
        pre.log_exposure = current_best_state.get("log_exposure")
        pre.rgb = current_best_state.get("observed_rgb")
        pre.density = current_best_state.get("density")
        pre.confidence_mask = current_best_state.get("confidence_mask")
        sep = type("sep", (object,), {"image_a": item.separation.image_a, "image_b": item.separation.image_b})()
        new_best_state = _build_best_state(
            pre, sep,
            current_best_state.get("film_stock"),
            current_best_state.get("physics_weight"),
            current_best_state.get("perceptual_weight"),
        )
    else:
        new_best_state = {}

    # Refresh status to show promoted as best
    # WP-4.1 Fix 3: strip any previous promote line so only one current line exists
    base = current_status.split("\n\n**Promoted to best:**")[0]
    promoted_status = base + f"\n\n**Promoted to best:** #{item.rank} `{item.candidate_id}` (loss={item.score.total_loss:.4f})"

    return pil_a, pil_b, recomb_pil, promoted_status, new_best_state


def identify_scenes_ui(best_state: dict | None, photo_context: str = "") -> tuple:
    """Ask the VLM to name the two scenes; fill the editable picker textboxes.

    Sends the working image to Anthropic (one small call), with the user's
    whole-photo description as ground truth when given. Soft-fails to a note
    without a key. The user can always type/edit the descriptions instead.
    """
    if not best_state:
        return "", "", "Run **Recover exposures** first, then identify the scenes."
    if not os.environ.get("ANTHROPIC_API_KEY"):
        return "", "", (
            "**Scene identification needs an `ANTHROPIC_API_KEY`** (Space secret / `.env`). "
            "You can still type the scene descriptions yourself below."
        )
    try:
        from app.restore import identify_scenes
        from app.demix import anthropic_vlm

        s1, s2 = identify_scenes(
            best_state["observed_rgb"], anthropic_vlm,
            context=(photo_context or "").strip() or None,
        )
        if not s1 and not s2:
            return "", "", "**Scene identification returned nothing** — type the descriptions yourself."
        note = (
            "Scenes identified — edit the descriptions if they're off, pick which one to "
            "recover, then click **AI Restore**. *(Image was sent to Anthropic for analysis.)*"
        )
        return s1, s2, note
    except Exception as exc:
        return "", "", f"**Scene identification failed:** {exc}"


def promote_alternate(evt: gr.SelectData, alternates: list | None, current_status: str = ""):
    """Click a take in the alternates gallery -> it becomes the restored main scene.

    WP-18 D2a: the status is re-rendered for the PROMOTED take — its faithfulness
    score and (when below threshold) a drift warning — so the text above the image
    always describes the pixels actually shown. Gallery select indices can arrive
    as (row, col) lists in some layouts; normalize to the flat index.
    """
    if not alternates or evt.index is None:
        return gr.update(), gr.update()
    raw = evt.index
    idx = int(raw[0]) if isinstance(raw, (list, tuple)) else int(raw)
    if idx < 0 or idx >= len(alternates):
        return gr.update(), gr.update()
    img, r = alternates[idx]

    from app.restore import DRIFT_R_THRESHOLD

    # One current "Showing take" line only (the promote_candidate strip pattern).
    base = (current_status or "").split("\n\n**Showing take")[0]
    line = f"\n\n**Showing take {idx + 1}** (faithfulness r={r:.2f})"
    if r < DRIFT_R_THRESHOLD:
        line += (
            " — ⚠️ this take drifted from your photo (the AI re-imagined rather "
            "than edited); treat it as an interpretation, not a restoration."
        )
    return to_pil(img), base + line


# NOT @GPU-decorated: this handler is network + CPU only (Replicate/Anthropic calls,
# scipy guided filter, numpy). Decorating it burned a ZeroGPU slice on API waits and
# imposed the 90s duration ceiling — best-of-3 plus one 429 backoff exceeded it, so
# ZeroGPU killed the handler mid-restore AFTER the paid API calls had gone out.
def restore_best_scene(
    best_state: dict | None,
    scene_choice: str = "Scene 1",
    scene_1: str = "",
    scene_2: str = "",
    photo_context: str = "",
    scribbles=None,
    tag_red: str = "", scene_red: str = "Scene 1",
    tag_orange: str = "", scene_orange: str = "Scene 1",
    tag_blue: str = "", scene_blue: str = "Scene 2",
    tag_magenta: str = "", scene_magenta: str = "Scene 2",
    upload: Image.Image | None = None,
    film_stock: str = "Portra 400",
    scan_type: str = "Auto",
    scan_calibration: str = "auto-exposed",
    auto_trim: bool = False,
    best_of_3: bool = True,
    ref_photo_1: list | None = None,  # file paths from gr.File (multiple)
    ref_photo_2: list | None = None,
    sam_objects: list | None = None,  # WP-22 tapped objects [{mask, tag, scene}]
    progress: "gr.Progress" = gr.Progress(),
) -> tuple:
    """Physics-anchored generative restoration (2026-07-16 pivot).

    Leads with ONE faithful photo of the scene the USER picked (human-in-the-loop:
    which scene matters is a preference, not something physics can rank). The other
    layer is offered as clearly-labelled best-effort. Optional user scribbles seed
    the physics split directly; per-color tags ("red = pool, Scene 1") additionally
    give the AI spatially grounded content anchors. Sends the image to Replicate
    only when a key is configured; soft-fails to a note otherwise.
    """
    # Standalone path: if the user clicked AI Restore without running Recover first,
    # preprocess the upload here so the densitometry the split needs exists. Recover
    # still adds ranked candidates + the confidence map; this just unblocks Restore.
    if not best_state:
        if upload is None:
            return None, None, "Upload a photo first, then click AI Restore.", None, None
        try:
            st, cal = _map_scan_params(scan_type, scan_calibration)
            pre = preprocess_negative(
                upload, stock=film_stock,
                scan_type=st,
                scan_calibration=cal,
                auto_trim=bool(auto_trim),
                mask_mode="slope",
            )
            best_state = _minimal_best_state(pre, film_stock)
        except Exception as exc:
            return None, None, f"**Could not prepare the photo:** {exc}", None, None

    try:
        from app.restore import recover

        vlm = None
        if os.environ.get("ANTHROPIC_API_KEY"):
            from app.demix import anthropic_vlm

            vlm = anthropic_vlm

        # The picked scene leads; the other conditions the residual restore.
        s1, s2 = (scene_1 or "").strip(), (scene_2 or "").strip()
        scene2_leads = str(scene_choice).strip().lower().startswith("scene 2")
        if scene2_leads:
            headline, other = s2, s1
        else:
            headline, other = s1, s2

        # Optional tagged scribbles: each brush color carries a scene + label; swap
        # scene-1/scene-2 into headline/other slots per the user's choice.
        seeds_headline = seeds_other = None
        hints_headline = hints_other = None
        markup_rgb = None
        legend_headline = legend_other = None
        scribble_warnings: list[str] = []
        obs = best_state["observed_rgb"]
        try:
            from app.scribble import (
                parse_tagged_scribbles,
                marks_unreadable,
                markup_and_legends,
            )

            assignments = {
                "red": {"scene": scene_red, "tag": tag_red},
                "orange": {"scene": scene_orange, "tag": tag_orange},
                "blue": {"scene": scene_blue, "tag": tag_blue},
                "magenta": {"scene": scene_magenta, "tag": tag_magenta},
            }
            seeds_1, seeds_2, hints_1, hints_2 = parse_tagged_scribbles(
                scribbles, obs.shape[:2], assignments,
                trim_bbox_frac=best_state.get("trim_bbox_frac"),
            )
            markup_0 = None
            legend_1 = legend_2 = ""
            if seeds_1.any() or seeds_2.any():
                # WP-19: the strokes also reach the editor as PIXELS — an annotated
                # copy of the frame plus per-scene color legends for the prompt.
                markup_0, legend_1, legend_2 = markup_and_legends(
                    scribbles, obs.shape[:2], assignments,
                    trim_bbox_frac=best_state.get("trim_bbox_frac"), rgb=obs,
                )
            # WP-22: tapped-object masks (click-to-segment) join the strokes —
            # same guidance channels, machine-precise edges and honest counts.
            from app.scribble import objects_guidance

            o1, o2, oh1, oh2, omk, ol1, ol2 = objects_guidance(
                sam_objects, obs.shape[:2],
                trim_bbox_frac=best_state.get("trim_bbox_frac"),
                base_rgb=markup_0 if markup_0 is not None else obs,
            )
            seeds_1, seeds_2 = seeds_1 | o1, seeds_2 | o2

            def _join(a, b):
                return "; ".join(x for x in (a or "", b or "") if x)

            hints_1, hints_2 = _join(hints_1, oh1), _join(hints_2, oh2)
            legend_1, legend_2 = _join(legend_1, ol1), _join(legend_2, ol2)
            if omk is not None:
                markup_0 = omk
            if seeds_1.any() or seeds_2.any():
                if scene2_leads:
                    seeds_headline, seeds_other = seeds_2, seeds_1
                    hints_headline, hints_other = hints_2, hints_1
                    legend_headline, legend_other = legend_2, legend_1
                else:
                    seeds_headline, seeds_other = seeds_1, seeds_2
                    hints_headline, hints_other = hints_1, hints_2
                    legend_headline, legend_other = legend_1, legend_2
                markup_rgb = markup_0
            elif marks_unreadable(scribbles, obs.shape[:2]):
                # WP-18 D1b: painted, but every pixel failed the color gate — say so
                # instead of silently pretending the user never marked anything.
                scribble_warnings.append(
                    "⚠️ Your scribbles could not be read (heavily blended colors) — "
                    "the restore ran WITHOUT them. Paint with less overlap between "
                    "different colors, or lower the brush opacity."
                )
        except Exception:
            pass  # scribbles are optional; never block the restore on parsing

        # WP-20/21 identity references: real photos of the same people/places per
        # scene (multiple files), mapped scene-1/scene-2 -> headline/other.
        def _ident_rgbs(files):
            out = []
            for f in files or []:
                path = getattr(f, "name", None) or (f if isinstance(f, str) else None)
                if not path:
                    continue
                try:
                    with Image.open(path) as im:
                        arr = np.asarray(im.convert("RGB"), np.float32) / 255.0
                    out.append(arr)
                except Exception:
                    continue
            return out or None

        ident_1, ident_2 = _ident_rgbs(ref_photo_1), _ident_rgbs(ref_photo_2)
        if scene2_leads:
            identity_headline, identity_other = ident_2, ident_1
        else:
            identity_headline, identity_other = ident_1, ident_2

        # Progress: coarse but honest — each generative/referee step ticks once.
        # Wrapped defensively: gr.Progress outside a live queue context must never
        # break the restore itself (direct calls, tests, API clients).
        est_total = (4 + 3) if best_of_3 else 3

        def _progress(desc: str, _c=[0]) -> None:
            _c[0] += 1
            try:
                progress(min(_c[0] / est_total, 0.95), desc=desc)
            except Exception:
                pass

        _progress("Warming up…")
        notes: list[str] = []
        notes.extend(scribble_warnings)
        result = recover(
            obs,
            best_state.get("h_total"),
            best_state.get("confidence_mask"),
            vlm=vlm,
            notes=notes,
            scene_headline=headline or None,
            scene_other=other or None,
            context=(photo_context or "").strip() or None,
            seeds_headline=seeds_headline,
            seeds_other=seeds_other,
            hints_headline=hints_headline,
            hints_other=hints_other,
            headline_is_secondary=scene2_leads,
            markup_rgb=markup_rgb,
            legend_headline=legend_headline,
            legend_other=legend_other,
            identity_headline=identity_headline,
            identity_other=identity_other,
            progress_cb=_progress,
            # The app's current best physics separation (incl. gallery promotes)
            # feeds the generative step as the subtraction anchor.
            best_pair=(
                (best_state["image_a"], best_state["image_b"])
                if best_state.get("image_a") is not None
                and best_state.get("image_b") is not None
                else None
            ),
            n_candidates=3 if best_of_3 else 1,
        )

        dom_pil = to_pil(result.dominant) if result.dominant is not None else None
        sec_pil = to_pil(result.second) if result.second is not None else None

        if result.dominant is None:
            if os.environ.get("REPLICATE_API_TOKEN", "").strip():
                head = (
                    "**AI restoration did not return an image** (API error or transient "
                    "failure — see the note below). The physics separation above is unaffected."
                )
            else:
                head = (
                    "**No AI restoration produced.** This step needs a `REPLICATE_API_TOKEN` "
                    "(set it as a Space secret / in `.env`). The physics separation above "
                    "still works fully offline."
                )
        else:
            picked = headline if headline else "the dominant scene"
            head = (
                "### Restored main scene\n"
                f"A clean, well-exposed photograph of **{picked}**, reconstructed "
                "from your negative. Geometry follows the film; exposure and drowned areas "
                "are AI-restored."
            )
            if sec_pil is not None:
                head += (
                    f"\n\n**Second scene (best-effort):** the weaker exposure, derived by "
                    f"the physics split. Roughly **{result.dreamed_frac:.0f}%** of the frame "
                    f"was not pinned down by the film (or your marks) and is **AI-imagined** — "
                    f"treat it as a plausible impression, not a faithful record."
                )
            else:
                # WP-18 D2b: name the actual reason class — a scan whose
                # densitometry failed can NEVER produce a second layer; only
                # blame rate-limits when the physics prerequisites existed.
                if best_state.get("h_total") is None:
                    head += (
                        "\n\n_Second scene not derivable for this scan: densitometry "
                        "failed (no exposure map), so there is no residual to restore. "
                        "The main-scene restore above is unaffected._"
                    )
                else:
                    head += (
                        "\n\n_Second scene not returned — the notes below say why (a "
                        "Replicate rate-limit is retried automatically; balances under $5 "
                        "are throttled to 1 request at a time)._"
                    )

        if result.api_contacted:
            head += f"\n\n*{ASYM_CONSENT}*"
        detail = "\n\n".join(f"- {n}" for n in notes if n)
        status = head + (f"\n\n{detail}" if detail else "")

        # Best-of-N alternates: gallery of (thumbnail, faithfulness caption) + the
        # full-res arrays in state so a click can promote one into the main slot.
        alt_gallery = [
            (to_pil(img), f"take {i + 1} · faithfulness r={r:.2f}")
            for i, (img, r) in enumerate(result.alternates)
        ] or None
        # WP-18 D2a: keep (image, r) pairs so promotion can re-render the drift line.
        alt_state = [(img, r) for img, r in result.alternates] or None
        return dom_pil, sec_pil, status, alt_gallery, alt_state

    except Exception as exc:
        tb = traceback.format_exc()
        return None, None, f"**Restore failed:** {exc}\n\n```\n{tb}\n```", None, None


def _sam_overlay(frame, objects, pending_mask=None, pending_pts=None):
    """Frame + committed fills (green/cyan) + in-progress mask and dots (yellow)."""
    from app.scribble import OBJECT_FILL_COLORS

    out = np.asarray(frame, np.float32) / 255.0
    for obj in objects or []:
        scene = "2" if str(obj.get("scene", "1")).endswith("2") else "1"
        col = np.asarray(OBJECT_FILL_COLORS[scene][1], np.float32) / 255.0
        m = np.asarray(obj["mask"], bool)
        out[m] = 0.5 * out[m] + 0.5 * col
    if pending_mask is not None and np.asarray(pending_mask).any():
        m = np.asarray(pending_mask, bool)
        col = np.asarray((255, 220, 0), np.float32) / 255.0
        out[m] = 0.45 * out[m] + 0.55 * col
    h, w = out.shape[:2]
    for x, y in pending_pts or []:  # dots so partial shapes stay visible
        x0, x1 = max(0, int(x) - 4), min(w, int(x) + 5)
        y0, y1 = max(0, int(y) - 4), min(h, int(y) + 5)
        out[y0:y1, x0:x1] = np.asarray((1.0, 0.86, 0.0), np.float32)
    return (np.clip(out, 0, 1) * 255).astype(np.uint8)


def _objects_md(objects) -> str:
    if not objects:
        return ""
    lines = [
        f"{i + 1}. {'🟢' if o.get('scene') == '1' else '🔵'} "
        f"**{o.get('tag') or 'object'}** — Scene {o.get('scene', '1')}"
        for i, o in enumerate(objects)
    ]
    return "**Tapped objects:** " + " · ".join(lines)


def _resolve_marks(frame, pts, mode):
    """(mask or None, note) for the current dots under the chosen mode (WP-23)."""
    if str(mode).lower().startswith("magic"):
        from app.segment import load_error, point_mask

        res = point_mask(np.asarray(frame, np.float32) / 255.0, [tuple(p) for p in pts])
        if res is None:
            err = load_error()
            return None, (
                "_Magic select unavailable"
                + (f" ({err[:120]})" if err else "")
                + " — switch to fill-from-dots or use the brush below._"
            )
        mask, iou = res
        return mask, (
            f"_Selected {100 * float(np.mean(mask)):.1f}% of the frame (confidence "
            f"{iou:.2f}). Tap again to refine, or **➕ Add object** to keep it._"
        )
    from app.scribble import polygon_mask

    mask = polygon_mask(pts, frame.shape[:2])
    if not mask.any():
        need = 3 - len(pts)
        return None, (
            f"_{len(pts)} dot{'s' if len(pts) != 1 else ''} placed — add "
            f"{max(need, 1)} more to close and fill the shape._"
        )
    return mask, (
        f"_Shape filled ({100 * float(np.mean(mask)):.1f}% of the frame). Add more "
        f"dots to refine the outline, or **➕ Add object** to keep it._"
    )


def sam_click(frame, pending_pts, objects, mode, evt: gr.SelectData):
    """One click = one dot; dots close a filled shape (default) or prompt SAM."""
    if frame is None:
        return gr.update(), pending_pts or [], None, "_Upload a photo first._"
    xy = evt.index
    pts = list(pending_pts or []) + [[float(xy[0]), float(xy[1])]]
    mask, note = _resolve_marks(frame, pts, mode)
    return _sam_overlay(frame, objects, mask, pts), pts, mask, note


def sam_add(frame, pending_pts, pending_mask, objects, tag, scene):
    """Commit the in-progress mask as a tagged object."""
    objects = list(objects or [])
    if pending_mask is None or not np.asarray(pending_mask).any():
        return (
            gr.update(), pending_pts or [], pending_mask, objects,
            _objects_md(objects), "_Tap an object first, then add it._", tag,
        )
    objects.append({
        "mask": np.asarray(pending_mask, bool),
        "tag": (tag or "").strip(),
        "scene": "2" if str(scene).endswith("2") else "1",
    })
    name = (tag or "").strip() or "object"
    return (
        _sam_overlay(frame, objects), [], None, objects, _objects_md(objects),
        f"_Added **{name}** ({len(objects)} object{'s' if len(objects) != 1 else ''}). "
        f"Tap the next one._", "",
    )


def sam_undo(frame, pending_pts, objects, mode="Fill shape from dots"):
    """Drop the last dot and recompute the in-progress shape/mask."""
    pts = list(pending_pts or [])[:-1]
    if frame is None or not pts:
        disp = _sam_overlay(frame, objects) if frame is not None else gr.update()
        return disp, [], None, "_Cleared the in-progress dots._"
    mask, _note = _resolve_marks(frame, pts, mode)
    return _sam_overlay(frame, objects, mask, pts), pts, mask, "_Removed the last dot._"


def sam_clear(frame):
    disp = _sam_overlay(frame, []) if frame is not None else None
    return disp, [], None, [], "", "_Cleared all tapped objects._"


def _working_frame_from_upload(img):
    """Upload -> working-geometry frame (EXIF transpose, megapixel guard, 1536 max).

    Mirrors preprocess_negative's intake exactly — the marking surfaces and the
    repair evidence must share the working image's geometry, or every mark lands
    misregistered against h_total/confidence_mask.
    """
    if img is None:
        return None
    from PIL import ImageOps
    from app.preprocessing import _guard_intake_size

    im = ImageOps.exif_transpose(img) or img
    im = _guard_intake_size(im).convert("RGB")
    w, h = im.size
    scale = min(1.0, 1536 / max(w, h))
    if scale < 1.0:
        im = im.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
    return np.asarray(im)


def _repair_source(target, main_pil, second_pil):
    """WP-24: the repair loop serves BOTH result images — flaws concentrate in
    the second scene (live referee scores ran 7/3 there)."""
    return second_pil if str(target).lower().startswith("second") else main_pil


def repair_click(target, main_pil, second_pil, repair_pts, evt: gr.SelectData):
    """Dot the flawed area directly on the restored result (WP-23)."""
    src = _repair_source(target, main_pil, second_pil)
    if src is None:
        return gr.update(), repair_pts or [], "_Run a restore first — then dot the area to fix._"
    arr = np.asarray(src.convert("RGB"))
    pts = list(repair_pts or []) + [[float(evt.index[0]), float(evt.index[1])]]
    from app.scribble import polygon_mask

    mask = polygon_mask(pts, arr.shape[:2])
    out = arr.astype(np.float32) / 255.0
    if mask.any():
        col = np.asarray((1.0, 0.25, 0.25), np.float32)
        out[mask] = 0.55 * out[mask] + 0.45 * col
    h, w = out.shape[:2]
    for x, y in pts:
        x0, x1 = max(0, int(x) - 4), min(w, int(x) + 5)
        y0, y1 = max(0, int(y) - 4), min(h, int(y) + 5)
        out[y0:y1, x0:x1] = np.asarray((1.0, 0.2, 0.2), np.float32)
    note = (
        f"_{len(pts)} dot{'s' if len(pts) != 1 else ''}"
        + (
            " — shape closed. Describe what belongs there and hit 🩹 Repair._"
            if mask.any()
            else f" — add {max(3 - len(pts), 1)} more to close the shape._"
        )
    )
    return (np.clip(out, 0, 1) * 255).astype(np.uint8), pts, note


def repair_apply(target, main_pil, second_pil, repair_pts, instruction, upload,
                 current_status: str = ""):
    """Re-render ONLY the dotted region; outside pixels return untouched."""
    src = _repair_source(target, main_pil, second_pil)
    if src is None:
        return gr.update(), gr.update(), gr.update(), "_Run a restore first._"
    if not repair_pts or len(repair_pts) < 3:
        return gr.update(), gr.update(), gr.update(), "_Dot a closed shape first (3+ dots)._"
    from app.restore import repair_region
    from app.scribble import polygon_mask

    base = np.asarray(src.convert("RGB"), np.float32) / 255.0
    mask = polygon_mask(repair_pts, base.shape[:2])
    observed = None
    if upload is not None:
        work = _working_frame_from_upload(upload)
        if work is not None:
            observed = np.asarray(work, np.float32) / 255.0
    notes: list[str] = []
    out, meta = repair_region(base, mask, instruction or "", observed, notes=notes)
    if out is None:
        why = "; ".join(notes) or "no result"
        return gr.update(), gr.update(), gr.update(), f"_Repair failed: {why}_"
    pil = to_pil(out)
    which = "second scene" if str(target).lower().startswith("second") else "main scene"
    status = (current_status or "") + (
        f"\n\n**🩹 Repaired the {which}** — “{(instruction or 'area cleanup').strip()[:80]}”: "
        + " ".join(notes)
    )
    if which == "second scene":
        return gr.update(), pil, status, "_Done — the second scene was updated. Dot another area to keep going._"
    return pil, gr.update(), status, "_Done — the main scene was updated. Dot another area to keep going._"


def finalize_current(current_main, current_status: str = "") -> tuple:
    """WP-21: re-render the CURRENTLY SHOWN main take at 2K (fidelity only).

    Operates on whatever is in the main slot — including a take the user promoted
    from the gallery — so approval and finalization compose. Failure keeps the
    1K image and explains in the status; success swaps the image in place.
    """
    from app.restore import finalize_take

    if current_main is None:
        return current_main, (current_status or "") + "\n\n_Nothing to finalize yet — run ✨ Restore first._"
    rgb = np.asarray(current_main.convert("RGB"), np.float32) / 255.0
    notes: list[str] = []
    out, _meta = finalize_take(rgb, notes=notes)
    extra = "".join(f"\n\n- {n}" for n in notes)
    if out is None:
        return current_main, (current_status or "") + "\n\n_2K finalize did not return an image:_" + extra
    pil = to_pil(out)
    line = f"\n\n**Finalized at {pil.width}×{pil.height}** — right-click / long-press the main image to save it."
    return pil, (current_status or "") + line + extra


def build_app() -> gr.Blocks:
    """Construct the Gradio Blocks application."""
    with gr.Blocks(
        title="Double Exposure Recovery",
        theme=gr.themes.Soft(),
    ) as demo:
        gr.Markdown(
            """
# 🎞️ Double Exposure Recovery
Two photos accidentally captured on one frame? **Upload the scan, click ✨ Restore,
and get your photo back.** Everything else is optional: describe or mark the two
scenes to guide the AI, add real photos of the people so faces come back faithful,
and open the physics workbench to see exactly what the film itself supports.
            """
        )
        gr.Markdown(TRANSPARENCY_NOTICE)

        with gr.Row():
            with gr.Column(scale=1):
                upload = gr.Image(
                    type="pil",
                    label="Negative scan",
                    sources=["upload"],
                )
                film_stock = gr.Dropdown(
                    choices=FILM_STOCKS,
                    value="Generic",
                    label="Film stock preset",
                )
                restore_btn_top = gr.Button("✨ Restore my photo", variant="primary")
                gr.Markdown(
                    "_One click does it all: the scenes are identified automatically "
                    "and the main one is recovered. Guide it on the right for better "
                    "results._"
                )
                with gr.Accordion("🔬 Physics workbench (advanced)", open=False):
                    num_candidates = gr.Slider(
                        minimum=1,
                        maximum=5,
                        value=3,
                        step=1,
                        label="Number of candidates",
                    )
                    include_deep_prior = gr.Checkbox(
                        value=False,
                        label="Deep prior separation — highest quality (slow, no API key needed)",
                        info=(
                            "Adds a Double-DIP candidate optimized per-image against the film "
                            "physics. On the 50-case 256px benchmark this is the strongest offline "
                            "source (beats the heuristic splits in 37/45 cases). Runs entirely "
                            "locally — no API key — but takes ~1-2 min per image, so it is opt-in."
                        ),
                    )
                    full_res_export = gr.Checkbox(
                        value=False,
                        label="Full-resolution export (slow, no API)",
                        info=(
                            "Post-process best separation to original scan resolution using "
                            "guided-filter upsampling of the exposure split ratio (grain from "
                            "your scan). Does not rerun optimization or API. Opt-in."
                        ),
                    )
                    scan_type = gr.Radio(
                        choices=["Auto", "Positive", "Negative"],
                        value="Auto",
                        label="Scan type",
                        info=(
                            "Auto: detect negative vs positive from luminance for display. "
                            "Positive: an already-inverted scan (lab JPEG / software-inverted) "
                            "of a double-exposed negative — physics is un-inverted internally; "
                            "true darkroom prints and slides are approximations. "
                            "Negative: force negative-scan densitometry (no polarity flip)."
                        ),
                    )
                    scan_calibration = gr.Radio(
                        choices=["auto-exposed", "calibrated-linear"],
                        value="auto-exposed",
                        label="Scanner",
                        info=(
                            "auto-exposed (default for real uploads): pin D_min to stock preset. "
                            "calibrated-linear: estimated white point (synthetic fixtures)."
                        ),
                    )
                    auto_trim = gr.Checkbox(
                        value=False,
                        label="Trim uniform border",
                        info="Opt-in: drop near-uniform rebate/letterbox edges (capped). Default off.",
                    )
                    asymmetric_recovery = gr.Checkbox(
                        value=False,
                        label=(
                            "Asymmetric recovery "
                            "(derive the second layer by physics subtraction; optional AI completion)"
                        ),
                        info=ASYM_CONSENT,
                    )
                    asymmetric_anchor = gr.Radio(
                        choices=["Auto", "Layer A", "Layer B"],
                        value="Auto",
                        label="Anchor",
                        info=(
                            "Which layer of the current best pair to treat as known. "
                            "Auto picks the larger mean exposure share."
                        ),
                    )
                    physics_weight = gr.Slider(
                        minimum=0.0,
                        maximum=5.0,
                        value=1.0,
                        step=0.1,
                        label="Physics loss weight",
                    )
                    perceptual_weight = gr.Slider(
                        minimum=0.0,
                        maximum=5.0,
                        value=0.5,
                        step=0.1,
                        label="Perceptual (LPIPS) weight",
                    )
                    run_btn = gr.Button("Recover exposures — physics separation", variant="secondary")

            with gr.Column(scale=2):
                gr.Markdown(
                    "## 🎨 Recover a clean photo\n"
                    "**✨ Restore my photo** works with nothing but an upload — the two "
                    "scenes are identified automatically. Everything here makes it "
                    "better: say what's in the frame, mark the scenes with the brush, "
                    "add real photos of the people, pick which scene to lead with. "
                    "Uses a generative editor (needs `REPLICATE_API_TOKEN`); the physics "
                    "separation in the workbench always works offline."
                )
                photo_context_box = gr.Textbox(
                    label="What's in this photo? (optional)",
                    placeholder="e.g. one shot is our backyard pool in California, the other is a museum in France",
                    lines=2,
                    info="Anything you know about the frame — it grounds both scene identification and the restore prompts.",
                )
                identify_btn = gr.Button("🔍 Identify the two scenes", variant="secondary")
                with gr.Row():
                    scene_1_box = gr.Textbox(
                        label="Scene 1", placeholder="e.g. a backyard swimming pool and patio",
                        lines=4, max_lines=10,
                    )
                    scene_2_box = gr.Textbox(
                        label="Scene 2", placeholder="e.g. an ornate gold picture frame on a wall",
                        lines=4, max_lines=10,
                    )
                scene_choice = gr.Radio(
                    ["Scene 1", "Scene 2"], value="Scene 1",
                    label="Which scene do you want recovered?",
                    info="You are the judge of which photo matters — the AI can't know.",
                )
                identify_note = gr.Markdown()
                with gr.Accordion(
                    "🖌️ Mark the scenes on the photo (optional — big quality boost)", open=False
                ):
                    gr.Markdown(
                        "**🎯 Dot around a thing, we fill the shape.** Click a few "
                        "dots around anything you recognize — even a faint ghost — "
                        "and the shape fills in like a paint bucket. Name it, pick "
                        "its scene, **➕ Add object**, move to the next. (Magic "
                        "select taps an object once and finds its edges "
                        "automatically — works best on clearly visible things.)"
                    )
                    sam_mode = gr.Radio(
                        ["Fill shape from dots", "Magic select (auto edges)"],
                        value="Fill shape from dots", label="Selection mode",
                    )
                    sam_display = gr.Image(
                        label="Click dots on the photo", type="numpy",
                        interactive=False,
                    )
                    with gr.Row():
                        sam_tag = gr.Textbox(
                            label="What is it?", placeholder="e.g. painting / pool / person",
                            scale=2,
                        )
                        sam_scene = gr.Radio(
                            ["Scene 1", "Scene 2"], value="Scene 1",
                            label="belongs to", scale=1,
                        )
                    with gr.Row():
                        sam_add_btn = gr.Button("➕ Add object", variant="secondary")
                        sam_undo_btn = gr.Button("↩️ Undo tap")
                        sam_clear_btn = gr.Button("🗑️ Start over")
                    sam_note = gr.Markdown()
                    sam_frame = gr.State(value=None)
                    sam_pending = gr.State(value=None)
                    sam_pending_mask = gr.State(value=None)
                    sam_objects_state = gr.State(value=None)
                    sam_summary = gr.Markdown()
                    gr.Markdown(
                        "**🖌️ Or paint freehand** — for regions that aren't neat "
                        "objects (skies, walls, halves of the frame):"
                    )
                    gr.Markdown(
                        "Paint over things you recognize, then (optionally) tag each "
                        "brush color below with **what it is** and **which scene it "
                        "belongs to** — e.g. paint the pool in red and tag red as "
                        "“pool”, Scene 1. The brushes are semi-transparent, so you can "
                        "see the photo underneath and layer colors over each other where "
                        "the two scenes overlap. Shading a whole shape works better than "
                        "outlining it. Your marks pin the physics split of the exposure "
                        "— the one thing the film alone cannot provide — and tagged marks "
                        "tell the AI what it's looking at. (The photo loads here as soon "
                        "as you upload it; you don't need to run Recover first.)"
                    )
                    scribble_canvas = gr.ImageEditor(
                        label="Brush colors: red · orange · blue · magenta (tag them below)",
                        type="numpy",
                        brush=gr.Brush(
                            colors=[
                                _PALETTE_RGBA["red"], _PALETTE_RGBA["orange"],
                                _PALETTE_RGBA["blue"], _PALETTE_RGBA["magenta"],
                            ],
                            default_color=_PALETTE_RGBA["red"],
                            default_size=12,
                        ),
                        transforms=(),
                        layers=False,
                        sources=(),
                    )
                    tag_boxes = {}
                    tag_scenes = {}
                    for _color, _default_scene in (
                        ("red", "Scene 1"), ("orange", "Scene 1"),
                        ("blue", "Scene 2"), ("magenta", "Scene 2"),
                    ):
                        with gr.Row():
                            tag_boxes[_color] = gr.Textbox(
                                label=f"{_color} strokes are…",
                                placeholder="e.g. pool / frame / person (optional)",
                                scale=2,
                            )
                            tag_scenes[_color] = gr.Radio(
                                ["Scene 1", "Scene 2"], value=_default_scene,
                                label="belongs to", scale=1,
                            )
                with gr.Accordion(
                    "🖼️ Real reference photos (optional) — bring faces back faithfully",
                    open=False,
                ):
                    gr.Markdown(
                        "A clear, ordinary photo of the **same people or place** in each "
                        "scene — another frame from the same roll works great. The AI uses "
                        "it for their true appearance (faces, hair, clothing) instead of "
                        "guessing; it never copies its pose or background."
                    )
                    with gr.Row():
                        ref_photo_1 = gr.File(
                            label="Scene 1 — people/place photos (up to 4)",
                            file_count="multiple", file_types=["image"],
                        )
                        ref_photo_2 = gr.File(
                            label="Scene 2 — people/place photos (up to 4)",
                            file_count="multiple", file_types=["image"],
                        )
                best_of_3 = gr.Checkbox(
                    value=True,
                    label="Best-of-3 (higher quality, 3× cost)",
                    info="Generative output varies run to run. Generate three takes of the "
                    "chosen scene, auto-rank them by faithfulness to your photo, and show "
                    "the best — the other takes appear below for you to pick from.",
                )
                restore_btn = gr.Button(
                    "🎨 AI Restore — recover the chosen scene", variant="primary"
                )
                with gr.Row():
                    restored_main = gr.Image(label="Restored main scene", type="pil")
                    restored_second = gr.Image(
                        label="Second scene (best-effort · AI-imagined)", type="pil"
                    )
                finalize_btn = gr.Button(
                    "🖼️ Finalize this photo at 2K (~$0.12) — sharper detail, same picture",
                    variant="secondary",
                )
                restore_status = gr.Markdown()
                alt_gallery = gr.Gallery(
                    label="All takes, most faithful first (click one to use it)",
                    columns=3, height=180, allow_preview=True,
                )
                alt_state = gr.State(value=None)
                with gr.Accordion(
                    "🩹 Fix an area — keep the rest pixel-identical", open=False
                ):
                    gr.Markdown(
                        "When a result is *almost* right: dot around the flawed "
                        "area on the image below, say what belongs there, and only "
                        "that region is re-rendered (~$0.08). Every pixel outside "
                        "your shape is carried over untouched — repairs can only "
                        "move the photo forward."
                    )
                    repair_target = gr.Radio(
                        ["Main scene", "Second scene"], value="Main scene",
                        label="Repair which image?",
                    )
                    repair_display = gr.Image(
                        label="Dot around the area to fix", type="numpy",
                        interactive=False,
                    )
                    repair_instruction = gr.Textbox(
                        label="What belongs there?",
                        placeholder="e.g. a second woman sitting in the green chair",
                    )
                    with gr.Row():
                        repair_btn = gr.Button("🩹 Repair this area", variant="secondary")
                        repair_reset_btn = gr.Button("Reset dots")
                    repair_note = gr.Markdown()
                    repair_pts = gr.State(value=None)

                with gr.Accordion(
                    "🔬 Physics engine room — what the film itself supports", open=True
                ):
                    gr.Markdown("### Comparison")
                    with gr.Row():
                        preprocessed_out = gr.Image(
                            label="Observed (preprocessed positive) + confidence overlay (if available)", type="pil"
                        )
                        recombined_out = gr.Image(
                            label="Best recombined (A ⊕ B)", type="pil"
                        )
                    gr.Markdown("### Best separation")
                    with gr.Row():
                        scene_a = gr.Image(label="Recovered scene A", type="pil")
                        scene_b = gr.Image(label="Recovered scene B", type="pil")
                    # WP-12: full-res exports (new outputs, do not replace working-res display)
                    gr.Markdown("### Full-resolution export (opt-in)")
                    with gr.Row():
                        full_a = gr.Image(label="Full-res scene A (download)", type="pil")
                        full_b = gr.Image(label="Full-res scene B (download)", type="pil")
                    status = gr.Markdown(label="Status & ranking")

                    # WP-4: gallery of all ranked + confidence map
                    gr.Markdown("### Ranked candidates (click to promote)")
                    candidates_gallery = gr.Gallery(
                        label="Click a pair to promote as best",
                        columns=3,
                        rows=2,
                        height="auto",
                        object_fit="contain",
                        show_label=True,
                        allow_preview=True,
                    )
                    conf_map_out = gr.Image(
                        label="Confidence map (grayscale: dark=toe, mid=valid, bright=shoulder)",
                        type="pil",
                    )

                with gr.Accordion(
                    "Advanced: physics-guided refinement (experimental)", open=False
                ):
                    gr.Markdown(
                        "Optionally refine the **best** result with gradient-based "
                        "optimization in VAE latent space, minimizing the same hybrid "
                        "physics + LPIPS loss. Slower than ranking; first run may "
                        "download VAE weights."
                    )
                    with gr.Row():
                        opt_steps = gr.Slider(
                            minimum=10,
                            maximum=150,
                            value=40,
                            step=5,
                            label="Optimization steps",
                        )
                        opt_lr = gr.Slider(
                            minimum=0.005,
                            maximum=0.2,
                            value=0.05,
                            step=0.005,
                            label="Learning rate",
                        )
                    enhance_btn = gr.Button(
                        "Enhance with Physics Optimization", variant="secondary"
                    )
                    with gr.Row():
                        refined_a_out = gr.Image(label="Refined scene A", type="pil")
                        refined_b_out = gr.Image(label="Refined scene B", type="pil")
                    refine_status = gr.Markdown()

        best_state = gr.State(value=None)
        ranked_state = gr.State(value=None)  # full list for gallery promote (WP-4)

        run_btn.click(
            fn=process_negative,
            inputs=[
                upload,
                film_stock,
                physics_weight,
                perceptual_weight,
                num_candidates,
                include_deep_prior,
                full_res_export,
                scan_type,
                scan_calibration,
                auto_trim,
                asymmetric_recovery,
                asymmetric_anchor,
            ],
            outputs=[
                preprocessed_out,
                recombined_out,
                scene_a,
                scene_b,
                status,
                best_state,
                ranked_state,
                candidates_gallery,
                conf_map_out,
                full_a,
                full_b,
            ],
        ).then(
            fn=lambda: (None, None, ""),
            inputs=None,
            outputs=[refined_a_out, refined_b_out, refine_status],
        ).then(
            # Sync the canvas to the exact working geometry after Recover, but ONLY
            # if it is still empty — never wipe strokes the user already painted
            # (the upload.change handler above normally fills it first). Sourced
            # from best_state (a gr.State) — wiring preprocessed_out as an input
            # would flip that display into an interactive upload widget.
            fn=lambda bs, canvas: (
                gr.update()
                if (isinstance(canvas, dict) and canvas.get("background") is not None)
                else (
                    (np.clip(bs["observed_rgb"], 0, 1) * 255).astype(np.uint8)
                    if bs else gr.update()
                )
            ),
            inputs=[best_state, scribble_canvas],
            outputs=[scribble_canvas],
        )

        enhance_btn.click(
            fn=enhance_best_result,
            inputs=[best_state, opt_steps, opt_lr],
            outputs=[refined_a_out, refined_b_out, refine_status],
        )

        identify_btn.click(
            fn=identify_scenes_ui,
            inputs=[best_state, photo_context_box],
            outputs=[scene_1_box, scene_2_box, identify_note],
        )

        restore_btn.click(
            fn=restore_best_scene,
            inputs=[
                best_state, scene_choice, scene_1_box, scene_2_box,
                photo_context_box, scribble_canvas,
                tag_boxes["red"], tag_scenes["red"],
                tag_boxes["orange"], tag_scenes["orange"],
                tag_boxes["blue"], tag_scenes["blue"],
                tag_boxes["magenta"], tag_scenes["magenta"],
                upload, film_stock, scan_type, scan_calibration, auto_trim,
                best_of_3, ref_photo_1, ref_photo_2, sam_objects_state,
            ],
            outputs=[restored_main, restored_second, restore_status, alt_gallery, alt_state],
        )

        # WP-21 one-click hero: same handler, same inputs — guidance boxes are
        # simply empty on the pure one-click path and the VLM fills the scenes.
        restore_btn_top.click(
            fn=restore_best_scene,
            inputs=[
                best_state, scene_choice, scene_1_box, scene_2_box,
                photo_context_box, scribble_canvas,
                tag_boxes["red"], tag_scenes["red"],
                tag_boxes["orange"], tag_scenes["orange"],
                tag_boxes["blue"], tag_scenes["blue"],
                tag_boxes["magenta"], tag_scenes["magenta"],
                upload, film_stock, scan_type, scan_calibration, auto_trim,
                best_of_3, ref_photo_1, ref_photo_2, sam_objects_state,
            ],
            outputs=[restored_main, restored_second, restore_status, alt_gallery, alt_state],
        )

        finalize_btn.click(
            fn=finalize_current,
            inputs=[restored_main, restore_status],
            outputs=[restored_main, restore_status],
        )

        alt_gallery.select(
            fn=promote_alternate,
            inputs=[alt_state, restore_status],
            outputs=[restored_main, restore_status],
        )

        # WP-23 repair loop: the repair surface mirrors whatever the main result
        # currently is (restore, promoted take, finalize, or a previous repair) —
        # .change fires on every programmatic update, so no handler needs to know.
        def _sync_repair(target, main_im, second_im):
            src = _repair_source(target, main_im, second_im)
            return (
                None if src is None else np.asarray(src.convert("RGB")), [], "",
            )

        for _trigger in (restored_main.change, restored_second.change, repair_target.change):
            _trigger(
                fn=_sync_repair,
                inputs=[repair_target, restored_main, restored_second],
                outputs=[repair_display, repair_pts, repair_note],
            )
        repair_display.select(
            fn=repair_click,
            inputs=[repair_target, restored_main, restored_second, repair_pts],
            outputs=[repair_display, repair_pts, repair_note],
        )
        repair_btn.click(
            fn=repair_apply,
            inputs=[repair_target, restored_main, restored_second, repair_pts,
                    repair_instruction, upload, restore_status],
            outputs=[restored_main, restored_second, restore_status, repair_note],
        )
        repair_reset_btn.click(
            fn=_sync_repair,
            inputs=[repair_target, restored_main, restored_second],
            outputs=[repair_display, repair_pts, repair_note],
        )

        # Load the marking canvas as soon as a photo is uploaded (same geometry as
        # preprocessing: EXIF transpose + max-side 1536), so users can mark scenes
        # before ever pressing Recover. auto_trim (opt-in) can shift geometry — the
        # post-Recover sync below covers that case without wiping existing strokes.
        _canvas_from_upload = _working_frame_from_upload

        upload.change(
            fn=_canvas_from_upload,
            inputs=[upload],
            outputs=[scribble_canvas],
        )

        # WP-22: the tap surface shows the same working-geometry frame as the
        # brush canvas; a new upload resets taps and objects.
        def _sam_from_upload(img):
            frame = _canvas_from_upload(img)
            return frame, frame, [], None, None, "", ""

        upload.change(
            fn=_sam_from_upload,
            inputs=[upload],
            outputs=[sam_frame, sam_display, sam_objects_state, sam_pending,
                     sam_pending_mask, sam_note, sam_summary],
        )
        sam_display.select(
            fn=sam_click,
            inputs=[sam_frame, sam_pending, sam_objects_state, sam_mode],
            outputs=[sam_display, sam_pending, sam_pending_mask, sam_note],
        )
        sam_add_btn.click(
            fn=sam_add,
            inputs=[sam_frame, sam_pending, sam_pending_mask, sam_objects_state,
                    sam_tag, sam_scene],
            outputs=[sam_display, sam_pending, sam_pending_mask, sam_objects_state,
                     sam_summary, sam_note, sam_tag],
        )
        sam_undo_btn.click(
            fn=sam_undo,
            inputs=[sam_frame, sam_pending, sam_objects_state, sam_mode],
            outputs=[sam_display, sam_pending, sam_pending_mask, sam_note],
        )
        sam_clear_btn.click(
            fn=sam_clear,
            inputs=[sam_frame],
            outputs=[sam_display, sam_pending, sam_pending_mask, sam_objects_state,
                     sam_summary, sam_note],
        )

        # WP-4: click in gallery promotes that candidate
        candidates_gallery.select(
            fn=promote_candidate,
            inputs=[ranked_state, status, best_state],
            outputs=[scene_a, scene_b, recombined_out, status, best_state],
        )

        gr.Markdown(
            """
---
**Setup:** Copy `.env.example` to `.env` and set `REPLICATE_API_TOKEN` for live
generative separation. Without it, the app runs in demo mode with heuristic candidates.
            """
        )

        # WP-4: warm-up LPIPS + VAE on first app load (progress indicator)
        demo.load(_warmup_models, inputs=None, outputs=None)

    return demo


def main() -> None:
    app = build_app()
    app.launch(
        server_name=os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1"),
        server_port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")),
        share=os.environ.get("GRADIO_SHARE", "").lower() in ("1", "true", "yes"),
    )


# --- WP-4 helpers (reused by process_negative + promote) ---

def _make_candidate_thumbnail(a: np.ndarray, b: np.ndarray, max_side: int = 160) -> Image.Image:
    """Create a small side-by-side thumbnail for gallery (A left, B right)."""
    pil_a = to_pil(a)
    pil_b = to_pil(b)
    # Resize preserving aspect (simple)
    def _resize(p: Image.Image, ms: int) -> Image.Image:
        w, h = p.size
        scale = min(1.0, ms / max(w, h))
        return p.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
    pil_a = _resize(pil_a, max_side)
    pil_b = _resize(pil_b, max_side)
    # Concat horizontal with thin separator
    sep = Image.new("RGB", (4, max(pil_a.height, pil_b.height)), (60, 60, 60))
    total_w = pil_a.width + sep.width + pil_b.width
    total_h = max(pil_a.height, pil_b.height)
    thumb = Image.new("RGB", (total_w, total_h), (30, 30, 30))
    thumb.paste(pil_a, (0, 0))
    thumb.paste(sep, (pil_a.width, 0))
    thumb.paste(pil_b, (pil_a.width + sep.width, 0))
    return thumb


def _confidence_mask_to_pil(mask: np.ndarray | None) -> Image.Image | None:
    """Grayscale visualization of WP-2 confidence mask (dark=TOE, mid=VALID, bright=SHOULDER)."""
    if mask is None:
        return None
    # Map 0->40, 1->128, 2->220 for visibility
    lut = np.array([40, 128, 220], dtype=np.uint8)
    gray = lut[mask.clip(0, 2)]
    return Image.fromarray(gray, mode="L").convert("RGB")


def _overlay_confidence(image_rgb: np.ndarray, mask: np.ndarray | None, alpha: float = 0.4) -> Image.Image:
    """Create observed image with grayscale confidence mask overlaid (for 'overlay' option)."""
    if mask is None:
        return to_pil(image_rgb)
    base = to_pil(image_rgb)
    conf = _confidence_mask_to_pil(mask)
    if conf is None:
        return base
    # Resize conf to match base if needed (should be same)
    if conf.size != base.size:
        conf = conf.resize(base.size, Image.Resampling.NEAREST)
    # Blend
    overlay = Image.blend(base.convert("RGBA"), conf.convert("RGBA"), alpha).convert("RGB")
    return overlay


# WP-4 warm-up (called on app load via demo.load)
def _warmup_models(progress: gr.Progress = gr.Progress()):
    """Force-load LPIPS (always) and VAE (best-effort) with visible progress."""
    progress(0.0, desc="Warming LPIPS (alex)...")
    try:
        from hybrid_loss import _get_shared_lpips
        _get_shared_lpips("alex")
    except Exception:
        pass
    progress(0.5, desc="Warming VAE (may download on first run or be unavailable in demo)...")
    try:
        from latent_optimizer import _get_shared_vae, DEFAULT_VAE_ID
        _get_shared_vae(DEFAULT_VAE_ID)
    except Exception:
        pass  # demo/offline ok
    progress(1.0, desc="Warm-up complete (models cached)")


def _minimal_best_state(preprocessed, film_stock: str) -> dict:
    """Restore-only best_state (WP-18 D4e): the subset of _build_best_state's
    schema the AI-Restore path consumes. Deliberately NO image_a/image_b — with
    no separation pair, recover() falls back to its heuristic anchor instead of
    treating the observed frame as a separation. Keep key names in lockstep with
    _build_best_state below.
    """
    return {
        "observed_rgb": preprocessed.rgb,
        "h_total": preprocessed.h_total,
        "confidence_mask": preprocessed.confidence_mask,
        "film_stock": film_stock,
        "trim_bbox_frac": getattr(preprocessed, "trim_bbox_frac", None),
    }


def _build_best_state(preprocessed, separation, film_stock, physics_weight, perceptual_weight):
    """Shared constructor for best_state dict to avoid schema drift between
    process_negative and promote_candidate (WP-4.1 Fix 4).
    """
    return {
        "image_a": separation.image_a,
        "image_b": separation.image_b,
        "log_exposure": preprocessed.log_exposure,
        "observed_rgb": preprocessed.rgb,
        "film_stock": film_stock,
        "physics_weight": physics_weight,
        "perceptual_weight": perceptual_weight,
        "density": preprocessed.density,
        "h_total": preprocessed.h_total,
        "confidence_mask": preprocessed.confidence_mask,
        "is_color": getattr(preprocessed, "is_color", False),
        # WP-18 D1c: scribble layers must be cropped by the same trim as the
        # working image before seeds are resized onto it.
        "trim_bbox_frac": getattr(preprocessed, "trim_bbox_frac", None),
    }


if __name__ == "__main__":
    main()