#!/usr/bin/env python3
"""Derive load-station ground truth from the raw CSV-parsed GT.

Two defects in the raw GT made the A1 numbers unreadable:

1. Double counting. 曾建程 and 賴彥碩 operate the SAME table -- 曾建程's rows even
   say "(與賴強化)". Their rows are two clerical records of one physical event,
   so the raw 33 "load cycles" in 09:30-10:30 are really ~17-19. 賴彥碩's times
   are minute-quantized, 曾建程's are second-precision, which is why they never
   collide exactly.
2. Wrong station. 阮文明 is the unload operator; his 2350 s work / 1021 s wait
   belong to a different station that this camera does not cover. Two of his
   rows also leaked into load_cycle_points because classify() keys on the
   substring 上片 without checking role.

Output keeps the raw file untouched and writes a station-scoped GT alongside it.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def cluster(points: list[dict], linkage: float) -> list[dict]:
    """Single-linkage merge of same-event rows recorded by different operators."""
    if not points:
        return []
    pts = sorted(points, key=lambda p: p["t"])
    groups: list[list[dict]] = [[pts[0]]]
    for p in pts[1:]:
        if p["t"] - groups[-1][-1]["t"] <= linkage:
            groups[-1].append(p)
        else:
            groups.append([p])
    merged = []
    for g in groups:
        # prefer a second-precision row for the timestamp; minute-quantized rows
        # carry +-30 s of clerical error we should not inherit.
        precise = [p for p in g if p["t"] % 60 != 0]
        t = min(p["t"] for p in precise) if precise else min(p["t"] for p in g)
        merged.append(
            {
                "t": int(t),
                "n_source_rows": len(g),
                "persons": sorted({p["person"] for p in g}),
                "t_quantized": not precise,
                "dur_sec": max(p.get("dur_sec", 0) for p in g),
                "action": g[0].get("action", ""),
            }
        )
    return merged


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--gt-json", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--station-role", default="load")
    ap.add_argument("--linkage", type=float, default=90.0)
    args = ap.parse_args()

    gt = json.loads(Path(args.gt_json).read_text(encoding="utf-8"))
    role_of = {e["person"]: e["role"] for e in gt["events"]}

    raw_pts = gt.get("load_cycle_points") or []
    kept = [p for p in raw_pts if role_of.get(p["person"]) == args.station_role]
    dropped_role = len(raw_pts) - len(kept)
    merged = cluster(kept, args.linkage)

    ivs = [i for i in gt.get("intervals") or [] if i.get("role") == args.station_role]
    work_sec = sum(i["t1"] - i["t0"] for i in ivs if i["kind"] == "work" and i.get("priority", 0) >= 2)
    wait_sec = sum(i["t1"] - i["t0"] for i in ivs if i["kind"] == "wait")

    out = {
        "derived_from": str(args.gt_json),
        "station_role": args.station_role,
        "linkage_sec": args.linkage,
        "video_t0_wall_sec": gt["video_t0_wall_sec"],
        "window_wall": gt["window_wall"],
        "raw_load_cycle_points": len(raw_pts),
        "dropped_wrong_role": dropped_role,
        "n_load_cycles": len(merged),
        "load_cycle_points": merged,
        "intervals": ivs,
        "work_person_sec": work_sec,
        "wait_person_sec": wait_sec,
        "operators": sorted({i["person"] for i in ivs}),
    }
    Path(args.out).write_text(json.dumps(out, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(
        f"raw={len(raw_pts)} -> dropped_role={dropped_role} -> kept={len(kept)} -> merged={len(merged)}\n"
        f"work_person_sec={work_sec} wait_person_sec={wait_sec} operators={out['operators']}"
    )
    return 0


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