#!/usr/bin/env python3
"""Single frame pass producing every pixel-derived feature Phase A1 needs.

Replaces the raw-Canny glass feature in run_phase_a1.py, which had a measured
CV of 0.024 over the hour -- the roller ROI was dominated by static structure
(furnace hood, white machine, the rollers themselves), so thresholding it was
thresholding noise.

The feature is a SATURATION DROP, not a grayscale difference. Glass on this line
is translucent: it lands on saturated red rollers and reads as a pale,
desaturated patch. Measured on visually labelled frames (7 glass / 3 empty):

    feature                       Cohen's d
    grayscale absdiff vs bg          0.33   <- what the baseline used
    value increase                   2.10
    saturation drop                  4.72   <- ranges do not overlap at all
    saturation drop AND value up     4.84

Grayscale absdiff scores ~0.66 whether the table is loaded or empty, because the
rollers are turning: it measures rotation, not glass.

Three things are computed here, all from one read of each JPG:

1. Background model  -- per-segment median saturation of the roller ROI, refined
   using only the quietest third of frames so a busy segment cannot bake glass
   into its own background.
2. Banded glass coverage -- desaturated fraction per band along the near->far
   (loading side -> furnace) axis, with person boxes masked out of both
   numerator and denominator. People lean over the rollers constantly; without
   masking they contaminate the glass signal.
3. Torso appearance descriptors -- HS histograms per detection, used by the
   offline tracklet stitcher.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import cv2
import numpy as np

SAT_DROP_THR = 25  # saturation units below background that count as glass
PERSON_PAD = 12  # px dilation on person boxes before masking

# Every clip opens with the camera's exposure/white-balance still settling, so
# the first frames come out effectively greyscale (measured leading runs: 20 /
# 0 / 54 / 36 frames on the four load clips). A greyscale frame has almost no
# saturation anywhere, which the glass rule reads as "the whole table just
# desaturated" -- coverage came out at 0.79 and manufactured a cycle at
# 08:00:00. Frames below this scene-level saturation are marked invalid rather
# than measured.
MIN_SCENE_SAT = 25


def norm_to_px(poly, w, h):
    return np.array([[p[0] * w, p[1] * h] for p in poly], dtype=np.float32)


def build_roi(zcfg, w, h, n_bands):
    """Return (mask, band_index, bbox) for the roller surface."""
    roi = zcfg["roller_glass_roi"]
    poly = norm_to_px(roi["polygon_norm"], w, h)
    mask_full = np.zeros((h, w), np.uint8)
    cv2.fillPoly(mask_full, [poly.astype(np.int32)], 255)

    x0, y0, bw, bh = cv2.boundingRect(poly.astype(np.int32))
    mask = mask_full[y0 : y0 + bh, x0 : x0 + bw] > 0

    near = norm_to_px(roi["near_edge_norm"], w, h).mean(axis=0)
    far = norm_to_px(roi["far_edge_norm"], w, h).mean(axis=0)
    axis = far - near
    denom = float(axis @ axis)

    yy, xx = np.mgrid[y0 : y0 + bh, x0 : x0 + bw]
    s = ((xx - near[0]) * axis[0] + (yy - near[1]) * axis[1]) / denom
    band = np.clip((np.clip(s, 0.0, 0.999) * n_bands).astype(np.int32), 0, n_bands - 1)
    return mask, band, (x0, y0, bw, bh)


def torso_descriptor(frame_hsv, bbox):
    x1, y1, x2, y2 = bbox
    bw, bh = x2 - x1, y2 - y1
    if bw < 8 or bh < 8:
        return np.zeros(64, np.float32)
    tx1 = int(x1 + 0.20 * bw)
    tx2 = int(x2 - 0.20 * bw)
    ty1 = int(y1 + 0.15 * bh)
    ty2 = int(y1 + 0.55 * bh)
    h, w = frame_hsv.shape[:2]
    tx1, tx2 = max(0, tx1), min(w, tx2)
    ty1, ty2 = max(0, ty1), min(h, ty2)
    if tx2 - tx1 < 4 or ty2 - ty1 < 4:
        return np.zeros(64, np.float32)
    patch = frame_hsv[ty1:ty2, tx1:tx2]
    hist = cv2.calcHist([patch], [0, 1], None, [8, 8], [0, 180, 0, 256])
    hist = hist.flatten().astype(np.float32)
    n = float(hist.sum())
    return hist / n if n > 0 else hist


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--frames-dir", required=True)
    ap.add_argument("--detections", required=True)
    ap.add_argument("--zones", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--segments", type=int, default=4)
    args = ap.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    zcfg = json.loads(Path(args.zones).read_text(encoding="utf-8"))
    n_bands = int(zcfg.get("tracker_rules", {}).get("n_bands", 8))

    frames = sorted(Path(args.frames_dir).glob("f_*.jpg"))
    dets_by_t = {}
    with Path(args.detections).open(encoding="utf-8") as fh:
        for line in fh:
            r = json.loads(line)
            dets_by_t[r["t_local"]] = r["detections"]
    n = len(frames)
    print(f"frames={n} detections_frames={len(dets_by_t)}", flush=True)

    probe = cv2.imread(str(frames[0]))
    h, w = probe.shape[:2]
    mask, band, (x0, y0, bw, bh) = build_roi(zcfg, w, h, n_bands)
    print(f"roi bbox=({x0},{y0},{bw},{bh}) px={int(mask.sum())} bands={n_bands}", flush=True)

    seg_bounds = [(i * n // args.segments, (i + 1) * n // args.segments) for i in range(args.segments)]

    cov = np.zeros(n, np.float32)
    band_cov = np.zeros((n, n_bands), np.float32)
    valid_frac = np.zeros(n, np.float32)
    greyscale = np.zeros(n, bool)
    desc_list: list[np.ndarray] = []
    desc_index: list[tuple[int, int]] = []  # (t_local, det_idx)
    bg_stack: list[np.ndarray] = []  # per-segment reference saturation, exported for the viewer

    for si, (a, b) in enumerate(seg_bounds):
        sats = np.zeros((b - a, bh, bw), np.int16)
        pmasks = np.zeros((b - a, bh, bw), bool)
        for i, t in enumerate(range(a, b)):
            img = cv2.imread(str(frames[t]))
            if img is None:
                continue
            hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
            sats[i] = hsv[y0 : y0 + bh, x0 : x0 + bw, 1].astype(np.int16)

            pm = np.zeros((bh, bw), np.uint8)
            for di, d in enumerate(dets_by_t.get(t, [])):
                bx1, by1, bx2, by2 = d["bbox_xyxy"]
                cv2.rectangle(
                    pm,
                    (bx1 - x0 - PERSON_PAD, by1 - y0 - PERSON_PAD),
                    (bx2 - x0 + PERSON_PAD, by2 - y0 + PERSON_PAD),
                    255,
                    -1,
                )
                desc_list.append(torso_descriptor(hsv, (bx1, by1, bx2, by2)))
                desc_index.append((t, di))
            pmasks[i] = pm > 0
        print(f"seg {si}: read {b-a} frames", flush=True)

        # greyscale frames must not be measured, and must not reach the
        # background model either
        usable = np.array([float(sats[i][mask].mean()) >= MIN_SCENE_SAT for i in range(b - a)])
        for i in np.flatnonzero(~usable):
            greyscale[a + int(i)] = True
        if (~usable).any():
            print(f"seg {si}: {int((~usable).sum())} greyscale frames excluded", flush=True)

        # pass 1 background, then refine on the quietest third so a busy segment
        # cannot bake its own glass into the reference saturation
        pool = np.flatnonzero(usable)
        if pool.size == 0:
            print(f"seg {si}: no usable frames", flush=True)
            bg_stack.append(np.zeros((bh, bw), np.uint8))
            del sats, pmasks
            continue
        bg = np.median(sats[pool], axis=0)
        rough = np.array(
            [float((((bg - sats[i]) > SAT_DROP_THR) & mask & ~pmasks[i]).sum()) for i in pool],
            np.float32,
        )
        quiet = pool[np.argsort(rough)[: max(8, pool.size // 3)]]
        bg = np.median(sats[quiet], axis=0)
        bg_stack.append(bg.astype(np.uint8))

        for i, t in enumerate(range(a, b)):
            if not usable[i]:
                continue  # coverage/band_cov stay 0, valid_frac stays 0
            valid = mask & ~pmasks[i]
            nv = int(valid.sum())
            valid_frac[t] = nv / float(mask.sum())
            if nv == 0:
                continue
            fg = ((bg - sats[i]) > SAT_DROP_THR) & valid
            cov[t] = fg.sum() / float(nv)
            for k in range(n_bands):
                sel = valid & (band == k)
                ns = int(sel.sum())
                band_cov[t, k] = (fg & sel).sum() / float(ns) if ns > 0 else 0.0
        del sats, pmasks
        print(f"seg {si}: cov mean={cov[a:b].mean():.4f} p95={np.percentile(cov[a:b],95):.4f}", flush=True)

    # coverage-weighted centroid along the near->far axis, in band units
    wsum = band_cov.sum(axis=1)
    idx = np.arange(n_bands, dtype=np.float32)
    centroid = np.where(wsum > 1e-6, (band_cov * idx).sum(axis=1) / np.maximum(wsum, 1e-6), np.nan)

    # The viewer recomputes the glass mask client-side to show exactly what the
    # detector decided, so it needs the same reference saturation. Segments are
    # stacked vertically into one grayscale PNG: row block i is segment i.
    cv2.imwrite(str(out_dir / "glass_bg.png"), np.vstack(bg_stack))

    np.savez_compressed(
        out_dir / "appearance.npz",
        desc=np.stack(desc_list) if desc_list else np.zeros((0, 64), np.float32),
        index=np.array(desc_index, np.int32) if desc_index else np.zeros((0, 2), np.int32),
    )
    (out_dir / "glass_bands.json").write_text(
        json.dumps(
            {
                "n_frames": n,
                "n_bands": n_bands,
                "sat_drop_thr": SAT_DROP_THR,
                "min_scene_sat": MIN_SCENE_SAT,
                "greyscale_frames": [int(v) for v in np.flatnonzero(greyscale)],
                "person_pad": PERSON_PAD,
                "roi_bbox": [x0, y0, bw, bh],
                "segments": [[int(a), int(b)] for a, b in seg_bounds],
                "bg_png": "glass_bg.png",
                "roi_polygon_norm": zcfg["roller_glass_roi"]["polygon_norm"],
                "coverage": [round(float(v), 5) for v in cov],
                "band_coverage": [[round(float(v), 5) for v in row] for row in band_cov],
                "centroid_band": [None if np.isnan(v) else round(float(v), 4) for v in centroid],
                "valid_frac": [round(float(v), 4) for v in valid_frac],
            },
            ensure_ascii=False,
        )
        + "\n",
        encoding="utf-8",
    )
    good = cov[cov > 0]
    print(
        f"DONE coverage mean={cov.mean():.4f} std={cov.std():.4f} CV={cov.std()/max(cov.mean(),1e-9):.3f} "
        f"p50={np.percentile(good,50):.4f} p95={np.percentile(good,95):.4f}",
        flush=True,
    )
    return 0


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