"""依頼 197: 元の記録 states.csv の実測角 (actual) の間隔を数える. 読むだけ.
python3 measure_gaps.py  → 同じフォルダに gaps.csv / runs.csv / by_cause.csv / by_stage.csv
"""
import csv, json, os, sys
from collections import defaultdict
import numpy as np

ROOT = "/Volumes/KIOXIA/r2r2r-data"
OUT = os.path.dirname(os.path.abspath(__file__))
SETS = ["teach_v1_trim2", "teach_pos1to4_deep_v1", "scripted_v1"]
LONG = 1.0

def phase_of(stage):
    s = (stage or "").upper()
    if "LIFT" in s or s == "STAGE6": return "carry_lift"
    if "RETREAT" in s or s.startswith("STAGE10"): return "retreat"
    if "APPROACH_PLACE" in s or s.startswith("STAGE7"): return "carry"
    if s in ("PICK",) or s.startswith("STAGE4") or s.startswith("STAGE5"): return "grasp"
    if "PLACE" in s or s.startswith("STAGE8") or s.startswith("STAGE9"): return "release"
    if "APPROACH_PICK" in s or s.startswith("STAGE3"): return "approach"
    if "HOME" in s or s in ("STAGE1", "STAGE2") or s.startswith("STAGE11") or s.startswith("STAGE12"): return "home"
    return "other:" + s

def load(run):
    p = f"{ROOT}/runs/run_pick_place/{run}/recording/states.csv"
    rows = list(csv.DictReader(open(p)))
    for r in rows: r["ts"] = float(r["ts"])
    return rows

def classify(rows, ds):
    """actual の 1 つの間ごとに, 間の中で起きたこと (原因) を決める."""
    acts = [i for i, r in enumerate(rows) if r["event_type"] == "actual" and r["j0"] != ""]
    out = []
    for a, b in zip(acts, acts[1:]):
        ra, rb = rows[a], rows[b]
        mid = rows[a + 1:b]
        types = [m["event_type"] for m in mid]
        # 直前の actual が stage 行のすぐあと (=送信前の読み) か
        prev = rows[a - 1]["event_type"] if a > 0 else ""
        if "gripper_cmd" in types or "gripper" in types:
            cause = "gripper"
        elif "commanded" in types:
            cause = "send"
        elif ds != "scripted_v1" and prev == "stage":
            cause = "send"          # record_teach_demo: stage → 送信前の読み → 送信 (wait=1.0)
        elif "stage" in types:
            cause = "stage_change"
        else:
            cause = "poll"
        out.append(dict(t0=ra["ts"], t1=rb["ts"], gap=rb["ts"] - ra["ts"], stage=ra["stage"],
                        phase=phase_of(ra["stage"]), cause=cause, between=";".join(types)))
    return out

def q(v, p): return float(np.percentile(v, p)) if len(v) else float("nan")

gap_rows, run_rows = [], []
for ds in SETS:
    idx = json.load(open(f"{ROOT}/datasets/mecharm_redblock_box_{ds}/index.json"))
    runs = [k for k, v in idx.items() if v.get("used")]
    for run in runs:
        g = classify(load(run), ds)
        for x in g: x.update(dataset=ds, run=run); gap_rows.append(x)
        v = np.array([x["gap"] for x in g])
        tot = v.sum()
        run_rows.append(dict(dataset=ds, run=run, n_actual=len(v) + 1, span_sec=round(tot, 3),
                             median=round(q(v, 50), 4), p90=round(q(v, 90), 4), max=round(v.max(), 3),
                             long_gt1s_count=int((v > LONG).sum()),
                             long_time_share=round(v[v > LONG].sum() / tot, 4),
                             mean_hz=round(len(v) / tot, 3)))

with open(f"{OUT}/gaps.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=["dataset", "run", "t0", "t1", "gap", "stage", "phase", "cause", "between"])
    w.writeheader(); [w.writerow({**x, "gap": round(x["gap"], 4)}) for x in gap_rows]
with open(f"{OUT}/runs.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=list(run_rows[0].keys())); w.writeheader(); w.writerows(run_rows)

def agg(key, fname):
    rows = []
    for ds in SETS:
        G = [x for x in gap_rows if x["dataset"] == ds]
        T = sum(x["gap"] for x in G)
        LT = sum(x["gap"] for x in G if x["gap"] > LONG)
        groups = defaultdict(list)
        for x in G: groups[x[key]].append(x["gap"])
        for k, v in sorted(groups.items()):
            v = np.array(v); lv = v[v > LONG]
            rows.append(dict(dataset=ds, **{key: k}, n=len(v), median=round(q(v, 50), 4), p90=round(q(v, 90), 4),
                             max=round(v.max(), 3), time_sec=round(v.sum(), 2), time_share=round(v.sum() / T, 4),
                             long_n=len(lv), long_time_sec=round(lv.sum(), 2),
                             share_of_long_time=round(lv.sum() / LT, 4) if LT else 0))
        v = np.array([x["gap"] for x in G])
        rows.append(dict(dataset=ds, **{key: "ALL"}, n=len(v), median=round(q(v, 50), 4), p90=round(q(v, 90), 4),
                         max=round(v.max(), 3), time_sec=round(T, 2), time_share=1.0,
                         long_n=int((v > LONG).sum()), long_time_sec=round(LT, 2), share_of_long_time=1.0))
    with open(f"{OUT}/{fname}", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys())); w.writeheader(); w.writerows(rows)
    return rows

for r in agg("cause", "by_cause.csv"): print(r)
print()
for r in agg("phase", "by_phase.csv"): print(r)