cv/snap.py

"""PAREIDOLIA deterministic feature snapping — the "code owns facts" half.

The VLM (MiniCPM-V) owns taste: it names the bolt that wants to be an eye and
hands over coarse normalized coordinates. This module owns facts: it snaps
each coarse point to the strongest nearby visual anchor using plain OpenCV on
CPU — Hough circles (eyes love circles: bolts, dials, knobs), good-features
corners, dark-blob centroids, and horizontally elongated edge segments for
mouths. No ML, no RNG: the same image + features always produce byte-identical
output, so a persisted menagerie record replays forever.

Contract (ARCHITECTURE.md §0 "Deterministic CV snapping", §2 step 3):

- search window = radius 12% of the image diagonal around each VLM point;
- candidates are scored by (anchor strength × proximity to the VLM point),
  weighted by how well the anchor kind suits the feature role;
- eyes prefer circle anchors, and the eye PAIR must stay coherent: if
  independent snapping skews the eye-line by >12° or changes the separation
  by >35% vs the VLM pair, the pair falls back to a rigid translation by the
  stronger anchor's delta (midpoint + angle of the VLM pair are preserved,
  the weaker eye is marked ``anchor_kind="pair"``);
- the mouth prefers a horizontally elongated contour/edge segment and is never
  snapped above the (snapped) eye midpoint;
- if no anchor beats ``MIN_SCORE`` the VLM point is kept verbatim
  (``anchor_kind="vlm"``) — the mist animation forgives ±15%;
- the mouth floor is enforced on the FINAL points regardless of anchor_kind:
  a kept-VLM mouth that sits at or above the eye midpoint is pushed just
  below it (``anchor_kind="vlm_corrected"``) so an inverted face can never
  ship.

Coordinate conventions: ``cx``/``cy`` are normalized to ``[0, 1]`` over
``width-1`` / ``height-1``; ``size`` is the feature's coarse diameter as a
fraction of the image diagonal (the ARCHITECTURE.md §3 schema); ``snap_delta``
is the distance moved, as a fraction of the image diagonal.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any, Optional

import cv2
import numpy as np

__all__ = ["snap_features", "overlay_debug", "MIN_SCORE", "SNAP_RADIUS_FRAC"]

# ----------------------------------------------------------------- constants

SNAP_RADIUS_FRAC = 0.12  # search radius as a fraction of the image diagonal
MIN_SCORE = 0.12  # below this, the VLM point wins (anchor_kind="vlm")

PAIR_MAX_SKEW_DEG = 12.0  # eye-line rotation tolerance vs the VLM pair
PAIR_MAX_SEP_CHANGE = 0.35  # eye separation ratio-change tolerance

# Corner response normalizer: cv2.cornerMinEigenVal output is pre-scaled by
# OpenCV — a crisp synthetic step corner peaks near ~0.06 (blockSize=5, uint8
# input, after our 5x5 blur) while blurred sensor noise stays ≤ ~0.016.
# Strength saturates at 1.0 so 0.08 keeps noise corners well under MIN_SCORE.
_CORNER_NORM = 0.08

# How well each anchor kind suits each feature role. Eyes love circles
# (bolts, dials, knobs); mouths love horizontal edge segments; corners are
# the weakest evidence everywhere (they fire on texture).
_KIND_WEIGHTS: dict[str, dict[str, float]] = {
    "eye": {"circle": 1.00, "blob": 0.70, "corner": 0.45, "edge": 0.30},
    "mouth": {"edge": 1.00, "blob": 0.75, "circle": 0.50, "corner": 0.35},
    "other": {"circle": 0.90, "blob": 0.80, "edge": 0.60, "corner": 0.50},
}

_MIN_ROI_SIDE = 12  # below this the window is too small to detect anything


@dataclass(frozen=True)
class _Anchor:
    """One candidate anchor in full-image pixel coordinates."""

    x: float
    y: float
    strength: float  # detector-specific, normalized to [0, 1]
    kind: str  # circle | corner | blob | edge


# ------------------------------------------------------------------ helpers


def _as_gray(image: np.ndarray) -> np.ndarray:
    """Accept BGR (contract) or already-gray uint8; return single-channel."""
    if image is None or not isinstance(image, np.ndarray) or image.ndim not in (2, 3):
        raise ValueError("snap_features expects an HxW or HxWx3 uint8 ndarray")
    if image.ndim == 3:
        return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    return image


def _clamp01(v: float) -> float:
    return min(1.0, max(0.0, float(v)))


def _role_class(role: Any) -> str:
    role = str(role or "")
    if role.startswith("eye"):
        return "eye"
    if role == "mouth":
        return "mouth"
    return "other"


def _to_px(f: dict, w: int, h: int) -> tuple[float, float]:
    """Normalized [0,1] feature coords → pixel coords (clamped on-image)."""
    return (
        _clamp01(f.get("cx", 0.5)) * (w - 1),
        _clamp01(f.get("cy", 0.5)) * (h - 1),
    )


def _proximity(d: float, radius: float) -> float:
    """Gaussian falloff: 1.0 at the VLM point, ~0.14 at the window edge."""
    return math.exp(-2.0 * (d / radius) ** 2)


# ---------------------------------------------------------------- detectors
# All detectors operate on a pre-blurred grayscale ROI and return anchors in
# ROI-local coordinates. Every parameter is fixed — nothing is sampled.


def _circle_candidates(roi: np.ndarray, want_r: float) -> list[_Anchor]:
    """HoughCircles, strength = disk-vs-annulus contrast × radius match."""
    if want_r > 0:
        min_r = max(2, int(round(want_r * 0.45)))
        max_r = max(min_r + 2, int(round(want_r * 1.9)))
    else:
        min_r, max_r = 3, max(6, int(min(roi.shape) * 0.45))
    circles = cv2.HoughCircles(
        roi,
        cv2.HOUGH_GRADIENT,
        dp=1.2,
        minDist=max(4.0, want_r if want_r > 0 else 8.0),
        param1=120,
        param2=18,
        minRadius=min_r,
        maxRadius=max_r,
    )
    if circles is None:
        return []
    out: list[_Anchor] = []
    for cx, cy, r in circles[0][:6]:
        contrast = _disk_contrast(roi, float(cx), float(cy), float(r))
        if want_r > 0:
            r_match = math.exp(-(((float(r) - want_r) / max(want_r, 1.0)) ** 2))
        else:
            r_match = 1.0
        strength = min(1.0, contrast * 1.6) * (0.4 + 0.6 * r_match)
        out.append(_Anchor(float(cx), float(cy), strength, "circle"))
    return out


def _disk_contrast(roi: np.ndarray, cx: float, cy: float, r: float) -> float:
    """|mean(disk interior) − mean(surrounding annulus)| / 255 — a bolt head
    or dial face separates from its plate; blurred noise does not."""
    h, w = roi.shape
    yy, xx = np.ogrid[:h, :w]
    d2 = (xx - cx) ** 2 + (yy - cy) ** 2
    inner = d2 <= (0.75 * r) ** 2
    ring = (d2 > (1.15 * r) ** 2) & (d2 <= (1.7 * r) ** 2)
    if int(inner.sum()) < 4 or int(ring.sum()) < 4:
        return 0.0
    return abs(float(roi[inner].mean()) - float(roi[ring].mean())) / 255.0


def _corner_candidates(roi: np.ndarray, want_r: float) -> list[_Anchor]:
    """goodFeaturesToTrack, strength from the min-eigenvalue response map."""
    corners = cv2.goodFeaturesToTrack(
        roi,
        maxCorners=10,
        qualityLevel=0.08,
        minDistance=max(4, int(want_r) if want_r > 0 else 6),
        blockSize=5,
    )
    if corners is None:
        return []
    response = cv2.cornerMinEigenVal(roi, blockSize=5)
    out: list[_Anchor] = []
    for pt in corners.reshape(-1, 2):
        x, y = float(pt[0]), float(pt[1])
        iy = min(response.shape[0] - 1, max(0, int(round(y))))
        ix = min(response.shape[1] - 1, max(0, int(round(x))))
        strength = min(1.0, float(response[iy, ix]) / _CORNER_NORM)
        out.append(_Anchor(x, y, strength, "corner"))
    return out


def _blob_candidates(roi: np.ndarray, want_r: float) -> list[_Anchor]:
    """Dark-blob centroids: adaptive threshold → contour moments, kept only
    when sized near the feature and actually darker than their surroundings."""
    side = min(roi.shape)
    block = int(round(want_r * 4)) | 1 if want_r > 0 else 21
    block = max(11, min(block, 51, (side - 1) | 1 if side > 2 else 3))
    if block < 3:
        return []
    binary = cv2.adaptiveThreshold(
        roi, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, block, 5
    )
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    want_area = math.pi * want_r * want_r if want_r > 0 else 0.0
    scored: list[tuple[float, Any]] = []
    for cnt in contours:
        area = float(cv2.contourArea(cnt))
        if area < 9.0:
            continue
        if want_area > 0 and not (0.15 * want_area <= area <= 6.0 * want_area):
            continue
        scored.append((abs(area - want_area), cnt))
    scored.sort(key=lambda t: t[0])
    out: list[_Anchor] = []
    for _, cnt in scored[:8]:
        m = cv2.moments(cnt)
        if m["m00"] <= 0:
            continue
        bx, by = m["m10"] / m["m00"], m["m01"] / m["m00"]
        mask = np.zeros(roi.shape, np.uint8)
        cv2.drawContours(mask, [cnt], -1, 255, -1)
        k = max(3, (int(want_r * 0.6) | 1) if want_r > 0 else 5)
        ring = cv2.dilate(mask, np.ones((k, k), np.uint8)) & ~mask
        if int((mask > 0).sum()) < 4 or int((ring > 0).sum()) < 4:
            continue
        darkness = (float(roi[ring > 0].mean()) - float(roi[mask > 0].mean())) / 255.0
        if darkness <= 0:
            continue  # spec: DARK blobs only
        d_eq = 2.0 * math.sqrt(float(cv2.contourArea(cnt)) / math.pi)
        if want_r > 0:
            size_match = math.exp(-0.5 * (((d_eq - 2 * want_r) / max(2 * want_r, 1.0)) ** 2))
        else:
            size_match = 1.0
        out.append(_Anchor(bx, by, min(1.0, darkness * 1.4) * size_match, "blob"))
    return out


def _edge_candidates(roi: np.ndarray, want_d: float) -> list[_Anchor]:
    """Horizontally elongated edge segments (mouths: slots, grilles, seams).

    Canny contours whose bounding box is clearly wider than tall. Strength =
    elongation × (soft) width match × vertical-gradient support × straightness.
    The last two terms separate a real seam (strong |dI/dy| along a near-
    straight run) from the wiggly low-contrast strings Canny traces on noise.
    """
    edges = cv2.Canny(roi, 60, 150)
    contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    if not contours:
        return []
    sobel_y = np.abs(cv2.Sobel(roi, cv2.CV_64F, 0, 1, ksize=3))
    out: list[_Anchor] = []
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if w < max(8, 0.3 * want_d) or w < 1.6 * max(h, 1):
            continue
        aspect = w / max(h, 1)
        elong = min(1.0, (aspect - 1.0) / 3.0)
        pts = cnt.reshape(-1, 2)
        ex, ey = float(pts[:, 0].mean()), float(pts[:, 1].mean())
        # a horizontal seam means strong vertical gradient along the contour
        grad_support = min(1.0, float(sobel_y[pts[:, 1], pts[:, 0]].mean()) / 450.0)
        # straight run ≈ arc 2w (traced out and back); wiggly noise is longer
        straightness = min(1.0, 2.2 * w / max(float(cv2.arcLength(cnt, False)), 1.0))
        if want_d > 0:
            size_match = math.exp(-0.5 * (((w - want_d) / max(want_d, 1.0)) ** 2))
        else:
            size_match = 1.0
        strength = elong * (0.5 + 0.5 * size_match) * grad_support * straightness
        out.append(_Anchor(ex, ey, strength, "edge"))
    out.sort(key=lambda a: (-a.strength, a.x, a.y))
    return out[:8]


# ------------------------------------------------------------- core snapping


def _gather_anchors(
    blurred: np.ndarray, px: float, py: float, radius: float, want_d: float
) -> list[_Anchor]:
    """Run all detectors on the search window; return full-image anchors."""
    h, w = blurred.shape
    x0 = max(0, int(math.floor(px - radius)))
    y0 = max(0, int(math.floor(py - radius)))
    x1 = min(w, int(math.ceil(px + radius)) + 1)
    y1 = min(h, int(math.ceil(py + radius)) + 1)
    roi = blurred[y0:y1, x0:x1]
    if min(roi.shape) < _MIN_ROI_SIDE:
        return []
    want_r = want_d / 2.0
    local = (
        _circle_candidates(roi, want_r)
        + _corner_candidates(roi, want_r)
        + _blob_candidates(roi, want_r)
        + _edge_candidates(roi, want_d)
    )
    return [_Anchor(a.x + x0, a.y + y0, a.strength, a.kind) for a in local]


def _best_anchor(
    anchors: list[_Anchor],
    px: float,
    py: float,
    radius: float,
    role_class: str,
    min_y: Optional[float],
) -> tuple[Optional[_Anchor], float]:
    """Pick the highest-scoring anchor within the search radius.

    ``min_y`` enforces the mouth rule: candidates at or above the snapped eye
    midpoint are discarded outright (a mouth is never above the eyes).
    Ties break on distance, then kind, then coordinates — fully deterministic.
    """
    weights = _KIND_WEIGHTS[role_class]
    best: Optional[_Anchor] = None
    best_key: tuple = ()
    best_score = 0.0
    for a in anchors:
        d = math.hypot(a.x - px, a.y - py)
        if d > radius:
            continue
        if min_y is not None and a.y <= min_y:
            continue
        score = weights[a.kind] * a.strength * _proximity(d, radius)
        key = (-score, d, a.kind, a.x, a.y)
        if best is None or key < best_key:
            best, best_key, best_score = a, key, score
    return best, best_score


def _snap_one(
    blurred: np.ndarray,
    feature: dict,
    radius: float,
    diag: float,
    min_y: Optional[float] = None,
) -> dict:
    """Snap a single feature; returns a NEW dict (inputs are never mutated)."""
    h, w = blurred.shape
    out = dict(feature)
    px, py = _to_px(feature, w, h)
    want_d = max(0.0, float(feature.get("size") or 0.0)) * diag
    role_class = _role_class(feature.get("role"))
    anchors = _gather_anchors(blurred, px, py, radius, want_d)
    best, score = _best_anchor(anchors, px, py, radius, role_class, min_y)
    if best is not None and score >= MIN_SCORE:
        out["cx"] = _clamp01(best.x / (w - 1)) if w > 1 else 0.0
        out["cy"] = _clamp01(best.y / (h - 1)) if h > 1 else 0.0
        out["snap_delta"] = math.hypot(best.x - px, best.y - py) / diag
        out["anchor_kind"] = best.kind
        out["anchor_score"] = round(score, 4)
    else:
        out["cx"] = _clamp01(feature.get("cx", 0.5))
        out["cy"] = _clamp01(feature.get("cy", 0.5))
        out["snap_delta"] = 0.0
        out["anchor_kind"] = "vlm"
        out["anchor_score"] = 0.0
    return out


def _enforce_pair_coherence(
    left: dict,
    right: dict,
    left_vlm: tuple[float, float],
    right_vlm: tuple[float, float],
    w: int,
    h: int,
    diag: float,
) -> None:
    """Keep the snapped eye pair geometrically honest (mutates result dicts).

    If independent snapping rotated the eye-line by more than
    ``PAIR_MAX_SKEW_DEG`` or stretched/shrunk the separation by more than
    ``PAIR_MAX_SEP_CHANGE`` vs the VLM pair, distrust the weaker anchor:
    rigidly translate the VLM pair by the STRONGER anchor's snap delta. The
    pair's midpoint offset and angle then match the VLM's intent, the stronger
    eye sits exactly on its anchor, and the weaker eye is re-derived
    (``anchor_kind="pair"``, ``anchor_score`` inherited from the evidence
    that placed it).
    """
    if left["anchor_kind"] == "vlm" and right["anchor_kind"] == "vlm":
        return  # nothing snapped, nothing to disagree about
    lvx, lvy = left_vlm
    rvx, rvy = right_vlm
    sep_v = math.hypot(rvx - lvx, rvy - lvy)
    if sep_v < 2.0:
        return  # degenerate VLM pair; geometry checks are meaningless
    lsx, lsy = left["cx"] * (w - 1), left["cy"] * (h - 1)
    rsx, rsy = right["cx"] * (w - 1), right["cy"] * (h - 1)
    sep_s = math.hypot(rsx - lsx, rsy - lsy)
    ang_v = math.degrees(math.atan2(rvy - lvy, rvx - lvx))
    ang_s = math.degrees(math.atan2(rsy - lsy, rsx - lsx))
    skew = abs((ang_s - ang_v + 180.0) % 360.0 - 180.0)
    sep_change = abs(sep_s / sep_v - 1.0)
    if skew <= PAIR_MAX_SKEW_DEG and sep_change <= PAIR_MAX_SEP_CHANGE:
        return
    # The eye whose anchor scored higher is trusted ('vlm' scores 0.0).
    if left["anchor_score"] >= right["anchor_score"]:
        strong, weak, strong_vlm, weak_vlm = left, right, (lvx, lvy), (rvx, rvy)
    else:
        strong, weak, strong_vlm, weak_vlm = right, left, (rvx, rvy), (lvx, lvy)
    dx = strong["cx"] * (w - 1) - strong_vlm[0]
    dy = strong["cy"] * (h - 1) - strong_vlm[1]
    weak["cx"] = _clamp01((weak_vlm[0] + dx) / (w - 1)) if w > 1 else 0.0
    weak["cy"] = _clamp01((weak_vlm[1] + dy) / (h - 1)) if h > 1 else 0.0
    weak["snap_delta"] = math.hypot(dx, dy) / diag
    weak["anchor_kind"] = "pair"
    weak["anchor_score"] = strong["anchor_score"]


def snap_features(image: np.ndarray, features: list[dict]) -> list[dict]:
    """Snap VLM feature points to the strongest nearby visual anchors.

    Args:
        image: HxWx3 BGR uint8 (HxW grayscale also accepted).
        features: dicts with at least ``cx``, ``cy`` (normalized [0,1]),
            ``size`` (coarse diameter / image diagonal; 0 or missing disables
            size matching) and ``role`` (``eye_left``/``eye_right``/``mouth``/
            anything else). Extra keys (``name``, …) pass through untouched.

    Returns:
        New dicts in input order with snapped ``cx``/``cy`` plus
        ``snap_delta`` (distance moved / image diagonal), ``anchor_kind``
        (``circle|corner|blob|edge|pair|vlm|vlm_corrected``) and
        ``anchor_score``.
        Inputs are never mutated. Fully deterministic — no RNG anywhere.
    """
    gray = _as_gray(image)
    h, w = gray.shape
    diag = math.hypot(w, h)
    radius = max(8.0, SNAP_RADIUS_FRAC * diag)
    blurred = cv2.GaussianBlur(gray, (5, 5), 1.2)

    results: list[Optional[dict]] = [None] * len(features)
    vlm_px = [_to_px(f, w, h) for f in features]

    # Pass 1 — eyes first: their final positions gate the mouth's floor.
    eye_idx = [i for i, f in enumerate(features) if _role_class(f.get("role")) == "eye"]
    for i in eye_idx:
        results[i] = _snap_one(blurred, features[i], radius, diag)

    left_i = next((i for i in eye_idx if features[i].get("role") == "eye_left"), None)
    right_i = next((i for i in eye_idx if features[i].get("role") == "eye_right"), None)
    if left_i is not None and right_i is not None:
        _enforce_pair_coherence(
            results[left_i], results[right_i], vlm_px[left_i], vlm_px[right_i], w, h, diag
        )

    eye_mid_y: Optional[float] = None
    if eye_idx:
        eye_mid_y = sum(results[i]["cy"] * (h - 1) for i in eye_idx) / len(eye_idx)

    # Pass 2 — everything else; mouths must land strictly below the eye line.
    for i, f in enumerate(features):
        if results[i] is not None:
            continue
        floor = eye_mid_y if _role_class(f.get("role")) == "mouth" else None
        results[i] = _snap_one(blurred, f, radius, diag, min_y=floor)

    # Pass 3 — enforce the mouth floor on the FINAL points regardless of
    # anchor_kind. Anchor candidates above the eye line were already filtered
    # in _best_anchor, but a kept-VLM mouth (anchor_kind="vlm") could still
    # sit above the snapped eye midpoint — an inverted face. Deterministic
    # correction: push it just below the eye line, preserving the VLM's own
    # eye-to-mouth scale.
    if eye_mid_y is not None:
        vlm_eye_mid_y = sum(vlm_px[i][1] for i in eye_idx) / len(eye_idx)
        for i, f in enumerate(features):
            if _role_class(f.get("role")) != "mouth":
                continue
            out = results[i]
            final_y = out["cy"] * (h - 1)
            if final_y > eye_mid_y:
                continue  # already strictly below the eye line
            mx, my = vlm_px[i]
            # 0.6 × the eye-to-mouth offset the VLM itself proposed (its
            # magnitude — the proposal may have been inverted), with a small
            # floor so an exactly-level mouth still moves strictly below.
            offset = 0.6 * abs(my - vlm_eye_mid_y)
            offset = max(offset, 0.02 * (h - 1), 1.0)
            new_y = min(eye_mid_y + offset, float(h - 1))  # clamp inside image
            out["cy"] = _clamp01(new_y / (h - 1)) if h > 1 else 0.0
            out["snap_delta"] = math.hypot(out["cx"] * (w - 1) - mx, new_y - my) / diag
            out["anchor_kind"] = "vlm_corrected"
    return results  # type: ignore[return-value]


# ------------------------------------------------------------- debug overlay

_COL_BEFORE = (0, 0, 255)  # red — the VLM's coarse guess
_COL_AFTER = (0, 255, 0)  # green — the snapped fact
_COL_WINDOW = (60, 60, 200)  # dim red — search radius
_COL_LINK = (0, 255, 255)  # yellow — the snap delta


def overlay_debug(
    image: np.ndarray, features_before: list[dict], features_after: list[dict]
) -> np.ndarray:
    """Render before(red)/after(green) points + search circles for eyeballing.

    Used by the eval agent's G2 grounding check (ARCHITECTURE.md §8): the red
    dot is what the VLM guessed, the dim red circle is the 12%-diagonal search
    window, the green dot is where deterministic code put the feature, with
    the winning ``anchor_kind`` labeled. Returns a new BGR uint8 image.
    """
    gray_or_bgr = image
    if gray_or_bgr.ndim == 2:
        canvas = cv2.cvtColor(gray_or_bgr, cv2.COLOR_GRAY2BGR)
    else:
        canvas = gray_or_bgr.copy()
    h, w = canvas.shape[:2]
    radius = int(round(max(8.0, SNAP_RADIUS_FRAC * math.hypot(w, h))))
    for before, after in zip(features_before, features_after):
        bx, by = _to_px(before, w, h)
        ax, ay = _to_px(after, w, h)
        b = (int(round(bx)), int(round(by)))
        a = (int(round(ax)), int(round(ay)))
        cv2.circle(canvas, b, radius, _COL_WINDOW, 1)
        cv2.line(canvas, b, a, _COL_LINK, 1)
        cv2.circle(canvas, b, 4, _COL_BEFORE, -1)
        cv2.circle(canvas, a, 4, _COL_AFTER, -1)
        kind = after.get("anchor_kind")
        if kind:
            cv2.putText(
                canvas,
                str(kind),
                (a[0] + 6, a[1] - 6),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.4,
                _COL_AFTER,
                1,
                cv2.LINE_AA,
            )
    return canvas