#!/usr/bin/env python3
"""Zones, occupancy and load cycles from re-tracked people + banded glass coverage.

Three defects in the baseline are fixed here.

Cycle state machine. run_phase_a1.py required the run to reach an `advancing`
state before it would emit anything, and silently dropped the cycle on
occupied -> empty. Any load whose push motion never cleared the P75 motion
threshold for 3 consecutive frames produced no output at all. Here `advance` is
recorded as an ATTRIBUTE of the cycle, never a precondition for emitting it.

Thresholds. The baseline set edge_thr to the 70th percentile of its own feature,
which structurally caps "occupied" at 30% of frames no matter what the line
actually did. Here the enter/exit thresholds come from an Otsu split of the
coverage histogram, with hysteresis so flicker around the split cannot spawn
cycles.

Work/wait. The baseline emitted one label per FRAME (any person in the work zone
made the whole frame "work"), so AI work seconds were capped at 3600 while GT
work is 5158 PERSON-seconds -- the gate was unreachable by construction. Here
everything is person-seconds, restricted to the load station, and work vs wait is
decided by per-track motion rather than by a wait polygon that covered the entire
working floor.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np


def norm_to_px(poly, w, h):
    return [[p[0] * w, p[1] * h] for p in poly]


def point_in_poly(x, y, poly):
    inside = False
    n = len(poly)
    j = n - 1
    for i in range(n):
        xi, yi = poly[i]
        xj, yj = poly[j]
        if ((yi > y) != (yj > y)) and (x < (xj - xi) * (y - yi) / (yj - yi + 1e-9) + xi):
            inside = not inside
        j = i
    return inside


def otsu(x: np.ndarray, bins: int = 64):
    hist, edges = np.histogram(x, bins=bins)
    hist = hist.astype(np.float64)
    total = hist.sum()
    if total == 0:
        return float(np.median(x))
    p = hist / total
    centres = (edges[:-1] + edges[1:]) / 2
    w0 = np.cumsum(p)
    m0 = np.cumsum(p * centres)
    mt = m0[-1]
    denom = w0 * (1 - w0)
    with np.errstate(divide="ignore", invalid="ignore"):
        var_b = np.where(denom > 1e-12, (mt * w0 - m0) ** 2 / np.maximum(denom, 1e-12), 0.0)
    return float(centres[int(np.argmax(var_b))])


def med_filt(x: np.ndarray, k: int) -> np.ndarray:
    if k <= 1:
        return x
    pad = k // 2
    xp = np.pad(x, pad, mode="edge")
    return np.array([np.median(xp[i : i + k]) for i in range(len(x))], np.float32)


def merge_runs(flags, min_len, max_gap):
    runs = []
    i, n = 0, len(flags)
    while i < n:
        if not flags[i]:
            i += 1
            continue
        j = i
        while j < n and flags[j]:
            j += 1
        runs.append([i, j])
        i = j
    if not runs:
        return []
    out = [runs[0]]
    for a, b in runs[1:]:
        if a - out[-1][1] <= max_gap:
            out[-1][1] = b
        else:
            out.append([a, b])
    return [(a, b) for a, b in out if b - a >= min_len]


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--tracks", required=True)
    ap.add_argument("--glass", required=True)
    ap.add_argument("--zones", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--video-t0", type=int, default=5400)
    args = ap.parse_args()

    out_dir = Path(args.out_dir)
    zcfg = json.loads(Path(args.zones).read_text(encoding="utf-8"))
    rules = zcfg["event_rules"]
    W, H = zcfg["image_size_hint"]

    zones_px = [(z["zone_id"], norm_to_px(z["polygon_norm"], W, H)) for z in zcfg["zones"]]
    station_px = norm_to_px(zcfg["load_station"]["polygon_norm"], W, H)

    glass = json.loads(Path(args.glass).read_text(encoding="utf-8"))
    cov = np.array(glass["coverage"], np.float32)
    centroid = np.array([np.nan if v is None else v for v in glass["centroid_band"]], np.float32)
    n = len(cov)

    frames = []
    with Path(args.tracks).open(encoding="utf-8") as fh:
        for line in fh:
            frames.append(json.loads(line))

    # ---- per-track speed, for the work/wait split -------------------------
    pos: dict[str, dict[int, tuple[float, float]]] = {}
    for fr in frames:
        for p in fr["people"]:
            x1, y1, x2, y2 = p["bbox_xyxy"]
            pos.setdefault(p["track_id"], {})[fr["t_local"]] = ((x1 + x2) / 2.0, float(y2))
    win = int(rules["idle_window_sec"])
    speed: dict[tuple[str, int], float] = {}
    for tid, series in pos.items():
        ts = sorted(series)
        for t in ts:
            near = [u for u in ts if abs(u - t) <= win]
            if len(near) < 2:
                speed[(tid, t)] = 0.0
                continue
            xs = np.array([series[u] for u in near], np.float32)
            span = max(near[-1] - near[0], 1)
            speed[(tid, t)] = float(np.linalg.norm(xs[-1] - xs[0]) / span)

    # ---- zones + activity, all in person-seconds --------------------------
    idle_thr = float(rules["idle_speed_thr_px_per_s"])
    work_zones = set(rules["work_zones"])
    sustain = int(rules["wait_min_sustain_sec"])
    raw: dict[str, dict[int, str]] = {}
    meta: dict[tuple[str, int], dict] = {}
    for fr in frames:
        t = fr["t_local"]
        for p in fr["people"]:
            x1, y1, x2, y2 = p["bbox_xyxy"]
            fx, fy = (x1 + x2) / 2.0, float(y2)
            zid = "other"
            for zname, poly in zones_px:
                if point_in_poly(fx, fy, poly):
                    zid = zname
                    break
            in_station = point_in_poly(fx, fy, station_px) or zid != "other"
            sp = speed.get((p["track_id"], t), 0.0)
            if not in_station:
                act = "offstation"
            elif zid in work_zones or sp > idle_thr:
                act = "work"
            else:
                act = "wait"
            raw.setdefault(p["track_id"], {})[t] = act
            meta[(p["track_id"], t)] = {"zone_id": zid, "speed": sp}

    # A pause is not a wait. The human chart only records 等待 when it lasts long
    # enough to be worth a row; brief stillness between load steps is still work.
    # Without this, momentary pauses at the cart inflate AI wait ~4x.
    for tid, series in raw.items():
        ts = sorted(series)
        i = 0
        while i < len(ts):
            if series[ts[i]] != "wait":
                i += 1
                continue
            j = i
            while j < len(ts) and series[ts[j]] == "wait" and ts[j] - ts[i] == j - i:
                j += 1
            if ts[j - 1] - ts[i] + 1 < sustain:
                for k in range(i, j):
                    series[ts[k]] = "work"
            i = j

    work_ps = sum(1 for s in raw.values() for v in s.values() if v == "work")
    wait_ps = sum(1 for s in raw.values() for v in s.values() if v == "wait")
    off_ps = sum(1 for s in raw.values() for v in s.values() if v == "offstation")

    out_frames = []
    for fr in frames:
        t = fr["t_local"]
        people = []
        for p in fr["people"]:
            tid = p["track_id"]
            m = meta[(tid, t)]
            people.append(
                {
                    "track_id": tid,
                    "bbox_xyxy": p["bbox_xyxy"],
                    "confidence": round(float(p.get("confidence", 0.0)), 3),
                    "zone_id": m["zone_id"],
                    "activity_state": raw[tid][t],
                    "speed_px_s": round(m["speed"], 1),
                }
            )
        out_frames.append({"t_local": t, "t_abs": fr["t_abs"], "people": people})

    # ---- cycle state machine ---------------------------------------------
    cov_s = med_filt(cov, 5)
    split = otsu(cov_s)
    enter_thr = split * float(rules["cycle_enter_frac"]) * 2.0
    exit_thr = split * float(rules["cycle_exit_frac"]) * 2.0
    min_occ = int(rules["cycle_min_occupy_sec"])
    min_gap = int(rules["cycle_min_gap_sec"])
    min_adv = float(rules["cycle_min_advance_bands"])

    occupied = np.zeros(n, bool)
    state = False
    for i in range(n):
        if state:
            state = cov_s[i] > exit_thr
        else:
            state = cov_s[i] > enter_thr
        occupied[i] = state

    runs = merge_runs(occupied, min_occ, 3)
    cycles = []
    last_start = -10_000
    for a, b in runs:
        if a - last_start < min_gap:
            continue
        seg = centroid[a:b]
        seg = seg[~np.isnan(seg)]
        adv = float(seg.max() - seg[0]) if len(seg) >= 2 else 0.0
        peak = int(a + int(np.argmax(cov_s[a:b])))
        cycles.append(
            {
                "t_start": int(a),
                "t_peak": peak,
                "t_end": int(b),
                "dur_sec": int(b - a),
                "peak_coverage": round(float(cov_s[a:b].max()), 4),
                "advance_bands": round(adv, 2),
                "advanced": bool(adv >= min_adv),
                "t_start_abs": args.video_t0 + int(a),
                "t_peak_abs": args.video_t0 + peak,
                "t_end_abs": args.video_t0 + int(b),
            }
        )
        last_start = a

    # ---- occupancy intervals (scene level, for the timeline) --------------
    scene_work = [any(p["activity_state"] == "work" for p in f["people"]) for f in out_frames]
    scene_wait = [
        (not w) and any(p["activity_state"] == "wait" for p in f["people"])
        for w, f in zip(scene_work, out_frames)
    ]
    work_iv = [
        {"t0": args.video_t0 + a, "t1": args.video_t0 + b, "kind": "work"}
        for a, b in merge_runs(scene_work, 3, 2)
    ]
    wait_iv = [
        {"t0": args.video_t0 + a, "t1": args.video_t0 + b, "kind": "wait"}
        for a, b in merge_runs(scene_wait, 3, 2)
    ]

    with (out_dir / "frames_v2.jsonl").open("w", encoding="utf-8") as fh:
        for i, f in enumerate(out_frames):
            f["glass"] = {
                "coverage": round(float(cov[i]), 4),
                "coverage_smooth": round(float(cov_s[i]), 4),
                "centroid_band": None if np.isnan(centroid[i]) else round(float(centroid[i]), 2),
                "occupied": bool(occupied[i]),
            }
            fh.write(json.dumps(f, ensure_ascii=False) + "\n")

    (out_dir / "load_cycles_v2.json").write_text(
        json.dumps(
            {
                "cycles": cycles,
                "occupied": [bool(v) for v in occupied],
                "enter_thr": round(enter_thr, 4),
                "exit_thr": round(exit_thr, 4),
                "otsu_split": round(split, 4),
            },
            ensure_ascii=False,
        )
        + "\n",
        encoding="utf-8",
    )
    (out_dir / "occupancy_v2.json").write_text(
        json.dumps({"work": work_iv, "wait": wait_iv}, ensure_ascii=False) + "\n", encoding="utf-8"
    )

    ids = sorted({p["track_id"] for f in out_frames for p in f["people"]})
    summary = {
        "video_t0": args.video_t0,
        "n_frames": n,
        "n_cycles": len(cycles),
        "n_cycles_advanced": sum(1 for c in cycles if c["advanced"]),
        "enter_thr": round(enter_thr, 4),
        "exit_thr": round(exit_thr, 4),
        "work_person_sec": work_ps,
        "wait_person_sec": wait_ps,
        "offstation_person_sec": off_ps,
        "unique_track_ids": ids,
        "n_unique_tracks": len(ids),
    }
    (out_dir / "summary_v2.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps({k: v for k, v in summary.items() if k != "unique_track_ids"}, indent=2))
    return 0


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