#!/usr/bin/env python3
"""生成相机偏移 sidecar：data/latent/arope_sidecar.pt（结构见 SPEC「文件契约」）。

每个 clip 存逐帧 / 逐 cell 的玩家屏幕投影位移（dx, dy）px，训练侧据此算世界 RoPE
和 known/new 掩码。数据来源只有 data/meta/*.jsonl + data/logs/<session>/：
meta 里有 video_t / session / clip_id（转场行有 src_a / src_b / n_pre），与 .pt 里的
元数据是同一份（encode_latents.py 直接从 meta 抄进去），所以不必 torch.load 22,782 个
.pt（每个 3 MB）；只对 meta 缺的 clip 才回退读 .pt。

增益修正：verify_sidecar.py 对视频实测后得到 gain_correction（分 x/y，含符号），
本脚本用 --gain_fit 把它乘进 offset_px / frame_offset_px / walk_speed_px_s，
sidecar 里存的是**已乘上**的值；M_px832 仍是日志里 M 的原始换算，方便回溯。

    python -m actionrope.sidecar                      # 未修正版（gain = [1, 1]）
    python -m actionrope.sidecar --gain_fit outputs/sidecar_check/gain_fit.json
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import sys
import time
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from actionrope import geometry as G  # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA = os.path.join(ROOT, "data")
VERSION = "arope_sidecar_v1"
SESSIONS = ("20260821_190601_787", "20260823_201942_753")
CLIP_SECONDS = (G.FRAMES - 1) / G.FPS          # 首帧到末帧 5.0 s
STRAIGHT_RT_MIN = 0.9                          # 直行 clip 的实时率下限（世界降速的不算）

# worker 侧的全局 tracks：fork 之后子进程直接继承，不用 pickle 传 2.3M 样本的数组
_TRACKS: dict = {}


def load_meta(meta_dir: str) -> dict[str, dict]:
    """读全部 meta 行，按 clip_id 索引。"""
    rows = {}
    for f in sorted(glob.glob(os.path.join(meta_dir, "*.jsonl"))):
        with open(f, encoding="utf-8") as fh:
            for line in fh:
                r = json.loads(line)
                rows[r["clip_id"]] = r
    return rows


def meta_from_pt(pool_dir: str, clip_id: str) -> dict:
    """meta 里没有的 clip 回退读 .pt 的元数据（不取 latent）。"""
    d = torch.load(os.path.join(pool_dir, clip_id + ".pt"), map_location="cpu", weights_only=False)
    d.pop("latent", None)
    d["clip_id"] = clip_id
    if "video_t" not in d and "video_t0" in d:
        d["video_t"] = d["video_t0"]
    return d


def _plain_offsets(job):
    """worker：一个普通 clip 的逐帧偏移 + valid。"""
    clip_id, session, video_t, max_dt = job
    tracks = _TRACKS[session]
    off, valid, reason = G.sample_frame_offsets(tracks, video_t, return_valid=True, max_sample_dt=max_dt)
    want_t = float(video_t) + np.arange(G.FRAMES) / G.FPS
    idx = G.nearest_log_index(tracks, want_t)
    max_log_dt = float(np.abs(tracks["vid_t"][idx] - want_t).max())
    return clip_id, np.nan_to_num(off).astype(np.float32), bool(valid), reason, max_log_dt


def compute_plain(rows: dict, workers: int, max_dt: float) -> dict[str, dict]:
    jobs = [(cid, r["session"], r["video_t"], max_dt) for cid, r in rows.items()
            if r.get("session") in SESSIONS]
    out = {}
    with ProcessPoolExecutor(workers) as ex:
        for cid, off, valid, reason, max_log_dt in ex.map(_plain_offsets, jobs, chunksize=256):
            out[cid] = dict(frame_offset_px=off, valid=valid, reason=reason,
                            max_log_dt=max_log_dt, kind="plain")
    return out


def compute_transitions(rows: dict, plain: dict) -> dict[str, dict]:
    """转场 clip：帧 < n_pre 抄 src_a，之后接 src_b 的增量（SPEC 的连续性规则）。"""
    out = {}
    for cid, r in rows.items():
        if r.get("session") != "transition":
            continue
        a, b = plain.get(r["src_a"]), plain.get(r["src_b"])
        if a is None or b is None:
            out[cid] = dict(frame_offset_px=np.zeros((G.FRAMES, 2), np.float32), valid=False,
                            reason="missing_src", max_log_dt=float("nan"), kind=r["kind"])
            continue
        off = G.transition_frame_offsets(a["frame_offset_px"], b["frame_offset_px"], int(r["n_pre"]))
        reasons = [x for x in (a["reason"], b["reason"]) if x]
        out[cid] = dict(frame_offset_px=off.astype(np.float32), valid=a["valid"] and b["valid"],
                        reason="src:" + "|".join(reasons) if reasons else "",
                        max_log_dt=max(a["max_log_dt"], b["max_log_dt"]), kind=r["kind"])
    return out


def is_straight(r: dict, rt_min: float = STRAIGHT_RT_MIN) -> bool:
    """直行 clip：21 个 cell 动作全同且 ≠ idle、实时率够高。"""
    acts = r.get("actions") or []
    return (len(acts) == G.CELLS and len(set(acts)) == 1 and acts[0] != 0
            and float(r.get("rt_ratio", 0.0)) >= rt_min)


def walk_speed_stats(rows: dict, clips: dict) -> dict:
    """直行 clip 的 |Δoffset|/秒：x 用左右走的（action 3/7），y 用上下走的（1/5）。
    按 session 分别给中位，整体取两 session 合并的中位。"""
    per = {}
    for cid, r in rows.items():
        c = clips.get(cid)
        if c is None or not c["valid"] or r.get("session") not in SESSIONS or not is_straight(r):
            continue
        act = r["actions"][0]
        d = c["frame_offset_px"][-1] - c["frame_offset_px"][0]
        s = per.setdefault(r["session"], dict(x=[], y=[], all=[]))
        if act in (3, 7):
            s["x"].append(abs(float(d[0])) / CLIP_SECONDS)
        if act in (1, 5):
            s["y"].append(abs(float(d[1])) / CLIP_SECONDS)
        s["all"].append(float(np.hypot(*d)) / CLIP_SECONDS)
    out = {}
    for sid, s in per.items():
        out[sid] = {k: (float(np.median(v)) if v else float("nan")) for k, v in s.items()}
        out[sid].update({f"n_{k}": len(v) for k, v in s.items()})
    allx = sum((s["x"] for s in per.values()), []); ally = sum((s["y"] for s in per.values()), [])
    out["median"] = [float(np.median(allx)) if allx else float("nan"),
                     float(np.median(ally)) if ally else float("nan")]
    return out


def build_sidecar(meta_dir: str, pool_dir: str, logs_dir: str, out_path: str, workers: int,
                  gain_fit: dict | None, max_dt: float, notes_extra: str = "") -> dict:
    t0 = time.time()
    rows = load_meta(meta_dir)
    pool_ids = sorted(f[:-3] for f in os.listdir(pool_dir) if f.endswith(".pt"))
    missing = [cid for cid in pool_ids if cid not in rows]
    if missing:
        print(f"meta 里缺 {len(missing)} 个 clip，回退读 .pt", flush=True)
        for cid in missing:
            rows[cid] = meta_from_pt(pool_dir, cid)
    rows = {cid: rows[cid] for cid in pool_ids}
    print(f"pool {len(pool_ids):,} 个 clip · meta {len(rows):,} 行 · 回退 {len(missing)}", flush=True)

    for sid in SESSIONS:
        _TRACKS[sid] = G.load_session_tracks(logs_dir, sid)
    M_px832 = _TRACKS[SESSIONS[0]]["M_px832"]
    for sid in SESSIONS[1:]:
        if not np.allclose(_TRACKS[sid]["M_px832"], M_px832):
            raise ValueError("两个 session 的 M 不一致，sidecar 的 meta.M_px832 无法唯一")

    plain = compute_plain(rows, workers, max_dt)
    trans = compute_transitions(rows, plain)
    clips_raw = {**plain, **trans}
    print(f"偏移计算完成 {len(clips_raw):,} 个 · {time.time() - t0:.1f}s", flush=True)

    # 增益修正（含符号）。先用未修正的值算步行速度，再乘 |gain|（速度是幅值，符号不参与）。
    gain = np.array((gain_fit or {}).get("gain_correction", [1.0, 1.0]), np.float64)
    speed_raw = walk_speed_stats(rows, clips_raw)
    walk_speed = [speed_raw["median"][0] * abs(gain[0]), speed_raw["median"][1] * abs(gain[1])]

    clips = {}
    for cid, c in clips_raw.items():
        fo = c["frame_offset_px"].astype(np.float64) * gain[None, :]
        clips[cid] = {
            "offset_px": torch.from_numpy(G.frames_to_cells(fo).astype(np.float32)),
            "frame_offset_px": torch.from_numpy(fo.astype(np.float32)),
            "valid": bool(c["valid"]),
            "kind": c["kind"],
            "max_log_dt": float(c["max_log_dt"]),
        }
    n_valid = sum(int(c["valid"]) for c in clips.values())

    notes = (f"frame_offset_px = 玩家屏幕投影位移（p_i − p_0，M_px832 = 日志 M_px1280 按 832/1280、480/720 缩放）"
             f"× gain_correction；背景相对屏幕的位移 = −frame_offset_px。"
             f"日志按 video_t=video_t0+i/16 最近邻采样（与 make_clips 一致）；"
             f"valid=False 当任一帧最近邻样本时间差 > {max_dt} s、坐标 NaN、所选日志样本 vt 倒退或相邻差 > max({G.MAX_LOG_JUMP_S}, 1/16+2×{max_dt}+0.1) s"
             f"（跨 offset 台阶）或单帧位移 > {G.MAX_STEP_PX} px；"
             f"转场 clip 的 valid = src_a 且 src_b 均 valid。max_log_dt 是该 clip 最近邻时间差的最大值，"
             f"可据此放宽 valid 而不必重算。walk_speed_px_s 由直行 clip（21 cell 动作全同≠idle、rt_ratio≥{STRAIGHT_RT_MIN}）"
             f"的 |Δoffset|/{CLIP_SECONDS:.0f}s 取中位，已乘 |gain_correction|。")
    if gain_fit is None:
        notes += " gain_correction=[1,1]：尚未经 verify_sidecar.py 实测校验。"
    else:
        notes += " " + gain_fit.get("decision", "")
    if notes_extra:
        notes += " " + notes_extra

    meta = {
        "version": VERSION,
        "M_px832": M_px832.tolist(),
        "gain_correction": [float(gain[0]), float(gain[1])],
        "gain_fit": (gain_fit or {}).get("gain_fit", {}),
        "walk_speed_px_s": [float(walk_speed[0]), float(walk_speed[1])],
        "walk_speed_px_s_raw": speed_raw,
        "n_clips": len(clips), "n_valid": n_valid,
        "max_sample_dt": float(max_dt),
        "notes": notes,
    }
    sidecar = {"meta": meta, "clips": clips}

    # 先写临时文件再原子替换，避免半截文件被训练侧读到
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    tmp = out_path + ".tmp"
    torch.save(sidecar, tmp)
    os.replace(tmp, out_path)
    print(f"写出 {out_path} · {os.path.getsize(out_path) / 1e6:.0f} MB · {time.time() - t0:.1f}s", flush=True)
    print_stats(rows, clips_raw, clips, meta)
    return sidecar


def print_stats(rows: dict, clips_raw: dict, clips: dict, meta: dict):
    n = len(clips)
    kinds = {}
    for c in clips.values():
        k = kinds.setdefault(c["kind"], [0, 0]); k[0] += 1; k[1] += int(c["valid"])
    print(f"\n== 统计  n_clips={n:,}  n_valid={meta['n_valid']:,} ({meta['n_valid'] / n * 100:.1f}%)")
    for k, (a, b) in sorted(kinds.items()):
        print(f"   {k:<8} {a:6,} · valid {b:6,} ({b / max(a, 1) * 100:.1f}%)")
    reasons = {}
    for c in clips_raw.values():
        if not c["valid"]:
            key = c["reason"].split(" ")[0]
            reasons[key] = reasons.get(key, 0) + 1
    print(f"   invalid 原因: {reasons}")
    dts = np.array([c["max_log_dt"] for c in clips.values() if np.isfinite(c["max_log_dt"])])
    print("   若放宽 max_sample_dt 的 valid 数: " + "  ".join(
        f"{th}s→{int((dts <= th).sum()) + sum(1 for c in clips.values() if not np.isfinite(c['max_log_dt']) and c['valid']):,}"
        for th in (0.15, 0.3, 0.5, 1.0, 1.5)))
    for sid in SESSIONS:
        ids = [cid for cid, r in rows.items() if r.get("session") == sid]
        v = sum(int(clips[c]["valid"]) for c in ids)
        print(f"   {sid}: {len(ids):,} clip · valid {v:,} ({v / max(len(ids), 1) * 100:.1f}%)")

    last = np.stack([c["offset_px"][G.CELLS - 1].numpy() for c in clips.values() if c["valid"]])
    mag = np.hypot(last[:, 0], last[:, 1])
    pct = [5, 25, 50, 75, 95, 99]
    print(f"   |offset_px[20]| 分位 {pct}: {np.percentile(mag, pct).round(1).tolist()}  (valid clip)")
    print(f"   |dx[20]| 分位: {np.percentile(np.abs(last[:, 0]), pct).round(1).tolist()}")
    print(f"   |dy[20]| 分位: {np.percentile(np.abs(last[:, 1]), pct).round(1).tolist()}")
    sp = meta["walk_speed_px_s_raw"]
    for sid in SESSIONS:
        if sid in sp:
            s = sp[sid]
            print(f"   步行速度中位(未修正) {sid}: x {s['x']:.1f} px/s (n={s['n_x']}) · "
                  f"y {s['y']:.1f} px/s (n={s['n_y']}) · |v| {s['all']:.1f} px/s (n={s['n_all']})")
    print(f"   walk_speed_px_s (已乘 |gain|={[abs(g) for g in meta['gain_correction']]}): "
          f"{[round(v, 2) for v in meta['walk_speed_px_s']]}")


def main():
    ap = argparse.ArgumentParser(description="生成 ARoPE 相机偏移 sidecar")
    ap.add_argument("--meta", default=os.path.join(DATA, "meta"))
    ap.add_argument("--pool", default=os.path.join(DATA, "latent", "pool"))
    ap.add_argument("--logs", default=os.path.join(DATA, "logs"))
    ap.add_argument("--out", default=os.path.join(DATA, "latent", "arope_sidecar.pt"))
    ap.add_argument("--workers", type=int, default=16)
    ap.add_argument("--gain_fit", default=None,
                    help="verify_sidecar.py 写出的 gain_fit.json；给了就把 gain_correction 乘进去")
    ap.add_argument("--max_sample_dt", type=float, default=G.MAX_SAMPLE_DT,
                    help="valid 判定：最近邻日志样本时间差上限（秒），SPEC 默认 0.15")
    ap.add_argument("--notes", default="", help="附加到 meta.notes 的说明")
    args = ap.parse_args()
    gain_fit = None
    if args.gain_fit:
        with open(args.gain_fit, encoding="utf-8") as fh:
            gain_fit = json.load(fh)
    build_sidecar(args.meta, args.pool, args.logs, args.out, args.workers, gain_fit,
                  args.max_sample_dt, args.notes)


if __name__ == "__main__":
    main()