app/scribble.py

"""Scribble-guided exposure split (2026-07-16 user-annotation experiment).

The physics fixes H_A + H_B = H_total per pixel but the per-pixel SPLIT is
unidentifiable (MASTERPLAN I.5 #3) — priors must disambiguate, and the WP-15 arc
proved no objective can. User scribbles ARE the missing information: strokes
marking "this is scene A / scene B" seed a split-ratio field w that is propagated
edge-aware across the frame (multi-scale iterated guided filter, reusing
app.fullres.guided_filter). Then H_A = w*H_total and H_B = (1-w)*H_total —
sum-exact by construction, non-negative, geometry-registered.

Physics honesty: w splits INTENSITY, not structure. In true overlap regions it
attenuates rather than unmixes — the w-weighted render is a scene-biased, faithful
ANCHOR that the generative restore finishes (validated live on photos 127/131:
the previously-unseparable 131 house layer came out clean).
"""

from __future__ import annotations

from typing import Optional, Tuple

import numpy as np
from PIL import Image

from app.fullres import guided_filter

# Tagged-stroke palette: each brush color is a SLOT the user can label ("pool")
# and assign to a scene. Colors chosen for mutual distance in RGB so nearest-color
# classification of anti-aliased stroke pixels is unambiguous.
PALETTE = {
    "red": (255, 0, 0),
    "orange": (255, 165, 0),
    "blue": (0, 0, 255),
    "magenta": (255, 0, 255),
}
PALETTE_HEX = {"red": "#FF0000", "orange": "#FFA500", "blue": "#0000FF", "magenta": "#FF00FF"}
# Semi-transparent brush strings for the UI: strokes let the photo show through and
# can be layered over each other to mark overlapping shapes in the two scenes.
BRUSH_ALPHA = 0.5
PALETTE_RGBA = {
    name: f"rgba({r},{g},{b},{BRUSH_ALPHA})" for name, (r, g, b) in PALETTE.items()
}
# Default scene assignment per slot (UI can override): warm colors -> scene 1, cool -> 2.
DEFAULT_SCENES = {"red": "1", "orange": "1", "blue": "2", "magenta": "2"}
# Painted-pixel gates (tuned for semi-transparent brushes). A pixel counts as a
# stroke if its layer alpha exceeds _ALPHA_MIN. Classification matches the pixel's
# color DIRECTION (cosine), which is invariant to the brush's alpha and to whether
# the editor sends straight or premultiplied RGB — pure red at 50% alpha and at
# 100% alpha point the same way. A pixel is only assigned when its direction aligns
# with a palette color above _COS_MIN; red+blue overlaps blend to an off-axis purple
# and are rejected (left contested) rather than misattributed.
_ALPHA_MIN = 30
_COS_MIN = 0.955


def _phi(rgb: np.ndarray) -> np.ndarray:
    from densitometry import phi_display

    return phi_display(rgb)


def _resize_f(arr: np.ndarray, wh: Tuple[int, int]) -> np.ndarray:
    im = Image.fromarray((np.clip(arr, 0, 1) * 65535).astype(np.uint16))
    return np.asarray(im.resize(wh, Image.BILINEAR), np.float32) / 65535.0


def _resize_b(mask: np.ndarray, wh: Tuple[int, int]) -> np.ndarray:
    im = Image.fromarray(mask.astype(np.uint8) * 255)
    return np.asarray(im.resize(wh, Image.NEAREST)) > 127


def propagate_w(
    rgb: np.ndarray,
    seeds_a: np.ndarray,
    seeds_b: np.ndarray,
    iters: int = 40,
    eps: float = 2e-4,
) -> np.ndarray:
    """Multi-scale edge-aware propagation of scribble seeds to a dense w in [0,1].

    Coarse-to-fine (64->512 px): at each scale, iterate guided filtering (guide =
    observed luminance) with the seeds re-clamped as boundary pins each pass.
    Coarse scales carry the seeds across the frame; fine scales snap w to edges.
    w=1 means the pixel's exposure belongs to scene A; w=0 to scene B.
    """
    g_full = _phi(np.asarray(rgb, np.float32)).astype(np.float32)
    h0, w0 = g_full.shape
    seeds_a = np.asarray(seeds_a, bool)
    seeds_b = np.asarray(seeds_b, bool)
    if seeds_a.shape != (h0, w0) or seeds_b.shape != (h0, w0):
        raise ValueError(f"seed masks {seeds_a.shape}/{seeds_b.shape} != image {(h0, w0)}")

    # Scales actually run for this image (dedup after clamping to the image size).
    sides: list[int] = []
    for side in (64, 128, 256, 512):
        side = min(side, max(h0, w0))
        if side not in sides:
            sides.append(side)
        if side == max(h0, w0):
            break

    w_prev: Optional[np.ndarray] = None
    for k, side in enumerate(sides):
        sc = side / max(h0, w0)
        sh, sw = max(2, int(h0 * sc)), max(2, int(w0 * sc))
        g = _resize_f(g_full, (sw, sh))
        sa = _resize_b(seeds_a, (sw, sh))
        sb = _resize_b(seeds_b, (sw, sh))
        # WP-18 D1a: both-scene pixels are contested — pin only EXCLUSIVE seeds
        # (previously sb was pinned last and silently won every overlap).
        sa_x, sb_x = sa & ~sb, sb & ~sa
        w = np.full(g.shape, 0.5, np.float32) if w_prev is None else _resize_f(w_prev, (sw, sh))
        radius = max(2, side // 16)
        for _ in range(iters):
            w[sa_x] = 1.0
            w[sb_x] = 0.0
            w = np.clip(guided_filter(g, w, radius=radius, eps=eps), 0.0, 1.0)
        # Sharpen only at the FINAL two scales: the box-filter diffusion shrinks w
        # toward 0.5 each pass, so a mild pointwise gain re-commits decided pixels —
        # but at coarse scales whichever seed family covers more area floods the
        # frame, and sharpening there locks that in (rich-get-richer, measured on
        # photo 127's border strokes). Position-keyed, so small images still sharpen.
        gain = 1.6 if k == len(sides) - 1 else (1.3 if k == len(sides) - 2 else 1.0)
        if gain != 1.0:
            w = np.clip(0.5 + gain * (w - 0.5), 0.0, 1.0)
        w_prev = w

    # Final: bring w to full image resolution with one edge-aware pass, re-pin seeds.
    assert w_prev is not None
    w_full = _resize_f(w_prev, (w0, h0))
    w_full = np.clip(guided_filter(g_full, w_full, radius=16, eps=eps), 0.0, 1.0)
    w_full[seeds_a & ~seeds_b] = 1.0
    w_full[seeds_b & ~seeds_a] = 0.0
    return w_full.astype(np.float32)


def split_by_scribbles(
    observed_rgb: np.ndarray,
    h_total: np.ndarray,
    confidence_mask: np.ndarray,
    seeds_a: np.ndarray,
    seeds_b: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Scribble-guided split: returns (layer_a_rgb, layer_b_rgb, w).

    H_A = w*H_total, H_B = (1-w)*H_total (sum-exact); each layer rendered to a
    positive on its own P99 scale with the scan's chroma carried (demix convention).
    """
    from app.demix import _render_h_to_positive

    observed_rgb = np.asarray(observed_rgb, np.float32)
    h_total = np.asarray(h_total, np.float32)
    w = propagate_w(observed_rgb, seeds_a, seeds_b)
    # SOFT split for rendering: clamp w away from 0/1 so no pixel is fully erased.
    # When the two exposures overlap spatially (photo 128: patio and paintings-wall
    # share most pixels), a hard w blacks-out/blows-out whole regions and the
    # generative restore re-invents them from text. A biased-but-complete render
    # keeps every scene's structure visible for the editor to REMOVE rather than
    # hallucinate. Sum-exactness is preserved (w_r + (1-w_r) = 1). The returned w
    # stays unclamped — it is the honest attribution field for disclosure.
    #
    # EXCEPTION (reliability push): pixels the user EXPLICITLY painted are a direct
    # statement of ownership — hard-attribute them (w = 1/0 exactly) so marked
    # regions separate cleanly instead of carrying a 12% ghost of the other scene.
    # WP-18 D1a: a pixel painted with BOTH scenes' colors is a statement that both
    # scenes live there — it stays CONTESTED (soft value), never won by either side.
    w_r = np.clip(w, 0.12, 0.88)
    sa = np.asarray(seeds_a, bool)
    sb = np.asarray(seeds_b, bool)
    contested = sa & sb
    w_r[sa & ~contested] = 1.0
    w_r[sb & ~contested] = 0.0
    h_a = (w_r * h_total).astype(np.float32)
    h_b = ((1.0 - w_r) * h_total).astype(np.float32)
    valid = np.asarray(confidence_mask) == 1

    def render(h: np.ndarray) -> np.ndarray:
        if valid.any() and np.any(h[valid] > 0):
            p99 = float(np.percentile(h[valid], 99))
        else:
            p99 = float(np.percentile(h, 99)) if h.size else 1.0
        return _render_h_to_positive(h, ref_rgb=observed_rgb, norm_scale=max(p99, 1e-8))

    return render(h_a), render(h_b), w


def _crop_layer_by_frac(arr: np.ndarray, bbox_frac) -> np.ndarray:
    """Crop a stroke layer by the fractional bbox the working image was trimmed to.

    WP-18 D1c: the marking canvas is filled from the UNTRIMMED upload; when
    "Trim uniform border" cropped the working image, the stroke layers must be
    cropped identically (the WP-12 full-res pattern) or every seed lands
    misregistered after the resize.
    """
    tf, bf, lf, rf = bbox_frac
    oh, ow = arr.shape[:2]
    return arr[int(round(tf * oh)):int(round(bf * oh)), int(round(lf * ow)):int(round(rf * ow))]


def _painted_color_masks(
    editor_value, target_hw: Tuple[int, int], trim_bbox_frac=None
) -> Tuple[dict, bool]:
    """Per-palette-color painted masks from a Gradio ImageEditor value.

    The editor returns {"background": ..., "layers": [RGBA, ...], "composite": ...}.
    Painted strokes live in the layers' alpha; each opaque pixel is classified to
    its NEAREST palette color (robust to anti-aliased stroke edges). Tolerant of
    PIL/ndarray layers and of a missing/empty value (returns empty masks).

    Returns (masks, any_painted) — ``any_painted`` is True when opaque stroke
    pixels existed at all, so callers can distinguish "user painted nothing" from
    "every painted pixel failed the color gate" (WP-18 D1b).
    """
    h, w = target_hw
    masks = {name: np.zeros((h, w), bool) for name in PALETTE}
    any_painted = False
    layers = (editor_value or {}).get("layers") if isinstance(editor_value, dict) else None
    if not layers:
        return masks, any_painted

    names = list(PALETTE)
    centers = np.array([PALETTE[n] for n in names], np.float32)  # (K, 3)
    centers_u = centers / (np.linalg.norm(centers, axis=1, keepdims=True) + 1e-8)
    for layer in layers:
        if layer is None:
            continue
        arr = np.asarray(layer)
        if arr.ndim != 3 or arr.shape[2] < 3:
            continue
        if trim_bbox_frac is not None:
            arr = _crop_layer_by_frac(arr, trim_bbox_frac)
        if arr.shape[:2] != (h, w):
            pil = Image.fromarray(arr.astype(np.uint8))
            arr = np.asarray(pil.resize((w, h), Image.NEAREST))
        rgbv = arr[..., :3].astype(np.float32)
        alpha = arr[..., 3] if arr.shape[2] >= 4 else np.full((h, w), 255, arr.dtype)
        painted = np.asarray(alpha) > _ALPHA_MIN
        if not painted.any():
            continue
        any_painted = True
        norm = np.linalg.norm(rgbv, axis=-1, keepdims=True)
        rgbu = rgbv / (norm + 1e-8)                                  # (h,w,3) unit vectors
        cos = rgbu @ centers_u.T                                     # (h,w,K) cosine sim
        nearest = np.argmax(cos, axis=-1)
        aligned = np.take_along_axis(cos, nearest[..., None], axis=-1)[..., 0] >= _COS_MIN
        painted &= (norm[..., 0] > 20.0)  # ignore near-black transparent-fringe pixels
        for k, name in enumerate(names):
            masks[name] |= painted & aligned & (nearest == k)
    return masks, any_painted


def region_phrase(mask: np.ndarray) -> str:
    """Coarse human/model-readable location of a stroke mask (thirds grid)."""
    ys, xs = np.nonzero(mask)
    if ys.size == 0:
        return ""
    h, w = mask.shape
    cy, cx = float(ys.mean()) / h, float(xs.mean()) / w
    row = ["top", "middle", "bottom"][min(2, int(cy * 3))]
    col = ["left", "center", "right"][min(2, int(cx * 3))]
    loc = "center" if (row, col) == ("middle", "center") else f"{row} {col}"
    return loc


def parse_tagged_scribbles(
    editor_value,
    target_hw: Tuple[int, int],
    assignments: Optional[dict] = None,
    trim_bbox_frac=None,
) -> Tuple[np.ndarray, np.ndarray, str, str]:
    """Tagged strokes -> (seeds_scene1, seeds_scene2, hints_scene1, hints_scene2).

    ``assignments`` maps palette color name -> {"scene": "1"|"2", "tag": str}
    (missing colors fall back to DEFAULT_SCENES with no tag). Every painted color
    contributes its mask to its scene's seeds; tagged colors additionally yield a
    text hint like "the pool (bottom left)" so the generative prompt knows what
    the user pointed at and where. ``trim_bbox_frac`` (WP-18 D1c) is the fractional
    crop applied to the working image by auto_trim; stroke layers are cropped
    identically before resizing so seeds stay registered.

    When the user painted strokes but NONE survived the color gate (heavy layered
    blending), ``marks_unreadable()`` reports it — callers should warn instead of
    silently ignoring the marks (WP-18 D1b).
    """
    assignments = assignments or {}
    masks, _any_painted = _painted_color_masks(editor_value, target_hw, trim_bbox_frac)
    h, w = target_hw
    seeds = {"1": np.zeros((h, w), bool), "2": np.zeros((h, w), bool)}
    hints: dict[str, list] = {"1": [], "2": []}
    for name, mask in masks.items():
        if not mask.any():
            continue
        a = assignments.get(name) or {}
        scene = str(a.get("scene") or DEFAULT_SCENES[name]).strip()
        scene = "2" if scene.endswith("2") else "1"
        seeds[scene] |= mask
        tag = str(a.get("tag") or "").strip()
        if tag:
            loc = region_phrase(mask)
            hints[scene].append(f"the {tag} ({loc})" if loc else f"the {tag}")
    return (
        seeds["1"],
        seeds["2"],
        "; ".join(hints["1"]),
        "; ".join(hints["2"]),
    )


def parse_editor_scribbles(
    editor_value, target_hw: Tuple[int, int]
) -> Tuple[np.ndarray, np.ndarray]:
    """Untagged red/blue parse (back-compat): red -> A/scene 1, blue -> B/scene 2."""
    seeds_1, seeds_2, _h1, _h2 = parse_tagged_scribbles(editor_value, target_hw, None)
    return seeds_1, seeds_2


def marks_unreadable(editor_value, target_hw: Tuple[int, int]) -> bool:
    """True when the user painted strokes but NO pixel survived the color gate.

    WP-18 D1b: distinguishes "painted nothing" (False) from "painted, but heavy
    layered blending pushed every pixel off the palette axes" (True) so the UI can
    say the marks could not be read instead of silently ignoring them.
    """
    masks, any_painted = _painted_color_masks(editor_value, target_hw)
    return any_painted and not any(m.any() for m in masks.values())


def render_markup(rgb: np.ndarray, masks: dict, alpha: float = 0.55) -> np.ndarray:
    """The observed frame with the user's strokes re-rendered as color overlays.

    WP-19: this is the "annotated copy" reference image for the evidence-bundle
    restore — the strokes reach the editor as PIXELS (positions and counts),
    which text hints cannot carry. Re-rendering from the parsed masks (instead of
    using the editor composite) guarantees the overlay shares the working image's
    geometry, including any auto-trim crop already applied to the masks.
    """
    out = np.asarray(rgb, np.float32).copy()
    for name, mask in masks.items():
        if mask is None or not np.asarray(mask).any():
            continue
        col = np.asarray(PALETTE[name], np.float32) / 255.0
        m = np.asarray(mask, bool)
        out[m] = (1.0 - alpha) * out[m] + alpha * col
    return np.clip(out, 0.0, 1.0)


def markup_and_legends(
    editor_value,
    target_hw: Tuple[int, int],
    assignments: Optional[dict] = None,
    trim_bbox_frac=None,
    rgb: Optional[np.ndarray] = None,
) -> Tuple[Optional[np.ndarray], str, str]:
    """Annotated-copy reference image + per-scene stroke legends (WP-19).

    Returns (markup_rgb or None, legend_scene1, legend_scene2). The two legend
    strings describe the same strokes from each scene's perspective ("belongs to
    THIS photo" vs "the other photo"), because each restore call recovers a
    different target. None markup when nothing readable was painted or ``rgb``
    is missing.
    """
    assignments = assignments or {}
    masks, _any = _painted_color_masks(editor_value, target_hw, trim_bbox_frac)
    painted = {n: m for n, m in masks.items() if m.any()}
    if not painted or rgb is None:
        return None, "", ""

    def _legend(own_scene: str) -> str:
        parts = []
        for name, _m in painted.items():
            a = assignments.get(name) or {}
            scene = str(a.get("scene") or DEFAULT_SCENES[name]).strip()
            scene = "2" if scene.endswith("2") else "1"
            tag = str(a.get("tag") or "").strip() or "content"
            side = "this photo" if scene == own_scene else "the other photo"
            parts.append(f"{name} strokes mark '{tag}' — belongs to {side}")
        return "; ".join(parts)

    return render_markup(rgb, painted), _legend("1"), _legend("2")


# ---------------------------------------------------------------------------
# WP-22 — tapped-object guidance (click-to-segment; app.segment supplies masks)
# ---------------------------------------------------------------------------
# Object fills use colors OUTSIDE the brush palette so the stroke color-gate
# classification is untouched and legends stay unambiguous when both are used.
OBJECT_FILL_COLORS = {"1": ("green", (0, 200, 0)), "2": ("cyan", (0, 210, 255))}


def objects_guidance(
    objects,
    target_hw: Tuple[int, int],
    trim_bbox_frac=None,
    base_rgb: Optional[np.ndarray] = None,
):
    """Tapped objects -> (seeds_1, seeds_2, hints_1, hints_2, markup, legend_1, legend_2).

    ``objects`` is a list of {"mask": bool array (click geometry), "tag": str,
    "scene": "1"|"2"} committed in the tap UI. Masks are cropped by the same
    auto-trim bbox as strokes, resized to ``target_hw``, and unioned per scene
    into physics seeds. Hints aggregate counts per tag ("4× painting"): the
    count-adherence signal WP-19 proved matters. ``markup`` is ``base_rgb``
    (pass the stroke markup to compose, or the plain frame) with green/cyan
    object fills; None when there are no objects or no base. Legends mirror
    markup_and_legends' per-scene perspective.
    """
    h, w = target_hw
    seeds = {"1": np.zeros((h, w), bool), "2": np.zeros((h, w), bool)}
    tags: dict[str, dict[str, int]] = {"1": {}, "2": {}}
    for obj in objects or []:
        mask = obj.get("mask")
        if mask is None or not np.asarray(mask).any():
            continue
        m = np.asarray(mask, bool)
        if trim_bbox_frac is not None:
            m = _crop_layer_by_frac(m.astype(np.float32), trim_bbox_frac) > 0.5
        if m.shape != (h, w):
            m = _resize_b(m, (w, h))
        scene = "2" if str(obj.get("scene", "1")).strip().endswith("2") else "1"
        seeds[scene] |= m
        tag = str(obj.get("tag") or "").strip() or "object"
        tags[scene][tag] = tags[scene].get(tag, 0) + 1

    def _hint(scene: str) -> str:
        parts = []
        for tag, n in tags[scene].items():
            loc = region_phrase(seeds[scene])
            head = f"{n}× {tag}" if n > 1 else f"the {tag}"
            parts.append(f"{head} ({loc})" if loc else head)
        return "; ".join(parts)

    def _legend(own_scene: str) -> str:
        parts = []
        for scene in ("1", "2"):
            if not tags[scene]:
                continue
            color = OBJECT_FILL_COLORS[scene][0]
            side = "this photo" if scene == own_scene else "the other photo"
            counted = ", ".join(
                (f"{n}× {t}" if n > 1 else t) for t, n in tags[scene].items()
            )
            parts.append(
                f"{color} shapes precisely outline tapped objects of {side}: {counted}"
            )
        return "; ".join(parts)

    markup = None
    if base_rgb is not None and (seeds["1"].any() or seeds["2"].any()):
        markup = np.asarray(base_rgb, np.float32).copy()
        for scene in ("1", "2"):
            if seeds[scene].any():
                col = np.asarray(OBJECT_FILL_COLORS[scene][1], np.float32) / 255.0
                m = seeds[scene]
                markup[m] = 0.45 * markup[m] + 0.55 * col
        markup = np.clip(markup, 0.0, 1.0)
    return seeds["1"], seeds["2"], _hint("1"), _hint("2"), markup, _legend("1"), _legend("2")


def polygon_mask(points, hw: Tuple[int, int]) -> np.ndarray:
    """WP-23: dots -> filled shape (the MS-Paint-bucket contract Eddie asked for).

    ``points`` are (x, y) pixel coords clicked in order; three or more close the
    polygon and fill it. The user supplies a handful of dots around anything they
    can see — including faint ghosts no segmenter can find — and the fill does
    the shading. Returns an all-False mask below 3 points.
    """
    from PIL import Image as PILImage, ImageDraw

    h, w = hw
    im = PILImage.new("L", (w, h), 0)
    if points is not None and len(points) >= 3:
        ImageDraw.Draw(im).polygon(
            [(float(x), float(y)) for x, y in points], fill=255
        )
    return np.asarray(im) > 127