"""依頼110 E: π0.5 の失敗12件・成功5件の, 爪を閉じたあとの持ち上げとつかみを記録から調べる (読むだけ).

入力 (読むだけ):
  runs/pi05_policy/{run}/log.jsonl・recording/states.csv・recording/top_frame_log.csv・{run}.json
  lab/twin/20261107_失敗の再生と置き場所/B/{selection.json, timeline.csv, timeline_metrics.csv}
出力 (このフォルダ):
  e_fk_fit.json         爪先の高さの近似式 (B の Unity の値に当てはめ)
  e_chunks.csv          閉じる前後のかたまり・補正ごとの値 (1行 = 1つの出来事)
  e_close.csv           閉じた瞬間の動き・送信無視・状態のずれ (1行 = 1試行)
  e_summary.csv         試行ごとのまとめと原因の候補
"""
import csv
import json
from pathlib import Path

import numpy as np

ROOT = Path("/Volumes/KIOXIA/r2r2r-data")
RUNS = ROOT / "runs" / "pi05_policy"
TW = ROOT / "lab" / "twin" / "20261107_失敗の再生と置き場所"
B = TW / "B"
OUT = TW / "E_つかみと持ち上げ"

J0_TURN_DEG = 8.0          # B と同じ: 閉じたときから j0 が 8° 回ったら「箱へ運び始めた」
LIFT_OUT_MM = 20.0         # π0.5 が「持ち上げを出した」とみる高さ (成功の持ち上げ 26〜36mm より少し下)
STILL_DEG_S = 2.0          # 閉じる前 0.3 秒の関節の速さがこれ未満なら「止まっていた」
GAP_DEG = 1.5              # 閉じた瞬間の, 最後に送った目標と届いた角度の差がこれを超えたら「着ききっていない」


# ---------------------------------------------------------------- 爪先の高さの近似式
def _feats(J):
    J = np.atleast_2d(np.asarray(J, dtype=float))
    j1, j2, j3, j4 = [np.radians(J[:, k]) for k in (1, 2, 3, 4)]
    b = j1 + j2
    return np.stack([np.ones_like(j1), np.sin(j1), np.cos(j1), np.sin(b), np.cos(b),
                     np.sin(b) * np.cos(j4), np.cos(b) * np.cos(j4),
                     np.sin(b) * np.sin(j4) * np.cos(j3), np.cos(b) * np.sin(j4) * np.cos(j3),
                     np.sin(j4) * np.sin(j3)], 1)


def fit_height():
    ang = {(r["run"], r["i"]): r for r in csv.DictReader(open(B / "timeline.csv"))}
    X, y = [], []
    for r in csv.DictReader(open(B / "timeline_metrics.csv")):
        a = ang[(r["run"], r["i"])]
        X.append([float(a[f"j{k}"]) for k in range(6)]); y.append(float(r["tip_h_mm"]))
    X, y = np.array(X), np.array(y)
    idx = np.random.default_rng(110).permutation(len(y)); tr, te = idx[:1500], idx[1500:]
    c, *_ = np.linalg.lstsq(_feats(X[tr]), y[tr], rcond=None)
    e = np.abs(_feats(X[te]) @ c - y[te])
    c, *_ = np.linalg.lstsq(_feats(X), y, rcond=None)
    info = {"n": int(len(y)), "holdout_n": int(len(te)), "holdout_abs_err_median_mm": round(float(np.median(e)), 2),
            "holdout_abs_err_p95_mm": round(float(np.percentile(e, 95)), 2),
            "holdout_abs_err_max_mm": round(float(e.max()), 2), "coef": [float(v) for v in c],
            "note": "B の Unity の爪先の高さ (マットから mm) を j1〜j4 の三角関数で当てはめた. 実機とのずれは測っていない"}
    json.dump(info, open(OUT / "e_fk_fit.json", "w"), ensure_ascii=False, indent=1)
    return lambda q: float((_feats(np.asarray(q, dtype=float)[:6]) @ c)[0]), info


# ---------------------------------------------------------------- 記録を読む
def load_states(run):
    act, sends, grips, stages = [], [], [], []
    for r in csv.DictReader(open(RUNS / run / "recording" / "states.csv")):
        t = float(r["ts"])
        if r["event_type"] == "actual" and r["j0"] and r["stage"] != "home":    # HOME に戻る動きは数えない
            act.append([t] + [float(r[f"j{k}"]) for k in range(6)])
        elif r["event_type"] == "send_angles":
            sends.append((t, np.array([float(r[f"j{k}"]) for k in range(6)]), r["stage"]))
        elif r["event_type"] == "gripper_cmd":
            grips.append((t, int(r["gripper_state"]), r["stage"]))
        elif r["event_type"] == "stage":
            stages.append((t, r["stage"]))
    return np.array(act), sends, grips, stages


def q_at(act, t):
    """t 以前でいちばん新しい届いた角度."""
    k = np.searchsorted(act[:, 0], t, side="right") - 1
    return None if k < 0 else act[k, 1:]


def load_log(run):
    return [json.loads(l) for l in open(RUNS / run / "log.jsonl")]


def frame_log(run):
    p = RUNS / run / "recording" / "top_frame_log.csv"
    return [(int(r["frame_idx"]), float(r["capture_ts"])) for r in csv.DictReader(open(p))] if p.exists() else []


def r1(v):
    return None if v is None else round(float(v), 1)


def main():
    H, fit = fit_height()
    sel = json.load(open(B / "selection.json"))
    runs = [(r, "failure") for r in sel["failures"]] + [(r, "success") for r in sel["successes"]]
    btab = {r["run"]: r for r in csv.DictReader(open(B / "b_table.csv"))}
    ev_rows, close_rows, sum_rows = [], [], []

    for run, outcome in runs:
        act, sends, grips, stages = load_states(run)
        log = load_log(run)
        tj = json.load(open(RUNS / run / f"{run}.json"))
        closes = [g for g in grips if g[1] == 1]
        if not closes:
            sum_rows.append({"run": run, "outcome": outcome, "cause": "記録なし",
                             "basis": "爪を閉じる命令の記録がない (log.jsonl も推論1回で送っていない)"})
            continue
        t_close, _, close_stage = closes[0]
        q_close = q_at(act, t_close)
        h_close = H(q_close)
        j0_close = q_close[0]
        opens_after = [g[0] for g in grips if g[1] == 0 and g[0] > t_close]
        t_open = opens_after[0] if opens_after else np.inf
        # j0 が 8° 回った時刻
        after = act[(act[:, 0] > t_close) & (act[:, 0] < t_open)]
        turn = after[np.abs(after[:, 1] - j0_close) >= J0_TURN_DEG]
        t_turn = turn[0, 0] if len(turn) else np.inf
        t_end_win = min(t_open, t_turn, act[-1, 0] + 1)
        win = act[(act[:, 0] >= t_close) & (act[:, 0] <= t_end_win)]
        lift_total = max(H(q) for q in win[:, 1:]) - h_close if len(win) else 0.0

        # ---- 出来事の並び (かたまり と 補正 と 巻き戻し)
        events = []
        chunk_recs = [d for d in log if "chunk" in d and d.get("type") is None]
        close_chunk = None
        for d in chunk_recs:
            if d.get("grip_cmd") == 1 and d.get("executed") and d.get("t_exec_end") \
                    and abs(d["t_exec_end"] - t_close) < 5:
                close_chunk = d["chunk"]
        for d in log:
            t = d.get("type")
            if t is None and "chunk" in d:
                events.append(("chunk", d.get("t_exec_start") or d["t"], d))
            elif t == "correction" and d.get("action") in ("accept", "cancel"):
                events.append(("correction", d.get("t_start", d.get("t")), d))
            elif t == "rewind":
                events.append(("rewind", d.get("t", 0), d))
        events.sort(key=lambda e: e[1])

        # 閉じたかたまりの1つ前から, 窓の終わりまで
        model_lift_raw, model_lift_clip, model_lift_sent, model_lift_act = [], [], [], []
        corr_lift = []
        state_gaps = []
        note_clamps = []
        first_corr_seen = False
        for kind, t0, d in events:
            if kind == "chunk":
                c = d["chunk"]
                if close_chunk is None or c < close_chunk - 1:
                    continue
                if d["t"] > t_end_win:
                    # 窓の終わりのあとの推論は, 出していた値だけ見る (送らない分)
                    pass
                rel = c - close_chunk
                acts = np.asarray(d["actions"], dtype=float)
                st = np.asarray(d["state"], dtype=float)[:6]
                ms = np.asarray(d["measured_state"], dtype=float)[:6]
                gap = np.abs(st - ms)
                vias = d.get("vias") or []
                raw_last = np.asarray(vias[-1]["raw"]) if vias else acts[min(d["used_step"], len(acts) - 1), :6]
                plan_last = np.asarray(vias[-1]["planned_target"]) if vias else np.asarray(d["target"])
                tgt_last = np.asarray(vias[-1]["target"]) if vias else np.asarray(d["target"])
                sent = d.get("sent_targets") or []
                sent_last = np.asarray(sent[-1]) if sent else None
                aa = [v.get("actual_after") for v in vias if v.get("actual_after")]
                act_last = np.asarray(aa[-1]) if aa else None
                if act_last is None and d.get("executed") and d.get("t_exec_end"):
                    act_last = q_at(act, d["t_exec_end"] + 0.3)
                h_raw_horizon = [H(a) for a in acts[:, :6]]
                notes = sorted({n for v in vias for n in (v.get("notes") or [])} | set(d.get("notes") or []))
                retried = [v["step"] for v in vias if v.get("retried")]
                failed = [v["step"] for v in vias if v.get("called") and v.get("success") is False]
                row = {
                    "run": run, "outcome": outcome, "kind": "chunk", "chunk": c, "rel_to_close": rel,
                    "executed": bool(d.get("executed")), "grip_cmd": d.get("grip_cmd"),
                    "after_window": bool(d["t"] > t_end_win),
                    "state_source": d.get("state_source"),
                    "state_meas_gap_max_deg": r1(gap.max()), "state_meas_gap_joint": f"j{int(np.argmax(gap))}",
                    "state_meas_gap_j1_deg": r1(st[1] - ms[1]), "state_meas_gap_j2_deg": r1(st[2] - ms[2]),
                    "h_state_mm": r1(H(st) - h_close), "h_measured_mm": r1(H(ms) - h_close),
                    "j1_measured": r1(ms[1]), "j1_raw_used": r1(raw_last[1]), "j1_clipped": r1(plan_last[1]),
                    "j1_sent": r1(None if sent_last is None else sent_last[1]),
                    "j1_actual_after": r1(None if act_last is None else act_last[1]),
                    "h_raw_used_mm": r1(H(raw_last) - h_close),
                    "h_raw_horizon_max_mm": r1(max(h_raw_horizon) - h_close),
                    "h_raw_horizon_last_mm": r1(h_raw_horizon[-1] - h_close),
                    "h_clipped_mm": r1(H(plan_last) - h_close),
                    "h_target_from_actual_mm": r1(H(tgt_last) - h_close),
                    "h_sent_mm": r1(None if sent_last is None else H(sent_last) - h_close),
                    "h_actual_after_mm": r1(None if act_last is None else H(act_last) - h_close),
                    "j0_raw_used_minus_close": r1(raw_last[0] - j0_close),
                    "j0_raw_horizon_last_minus_close": r1(acts[-1, 0] - j0_close),
                    "notes": " / ".join(notes), "via_retried_steps": " ".join(map(str, retried)),
                    "via_failed_steps": " ".join(map(str, failed)),
                    "used_step": d.get("used_step"),
                }
                ev_rows.append(row)
                if rel in (-1, 0, 1, 2):
                    state_gaps.append((rel, d.get("state_source"), gap.max(), st[1] - ms[1]))
                if rel >= 1 and not first_corr_seen and d["t"] <= t_end_win:
                    model_lift_raw.append(max(h_raw_horizon) - h_close)
                    model_lift_clip.append(H(plan_last) - h_close)
                    if d.get("executed") and sent_last is not None:
                        model_lift_sent.append(H(sent_last) - h_close)
                        if act_last is not None:
                            model_lift_act.append(H(act_last) - h_close)
                    if any("クランプ" in n or "制限" in n for n in notes):
                        note_clamps.append((c, notes))
            elif kind in ("correction", "rewind"):
                if t0 is None or t0 < t_close or t0 > t_end_win:
                    continue
                first_corr_seen = True
                sp = d.get("start_pose"); ma = d.get("measured_after")
                hl = (H(ma) - H(sp)) if (sp and ma) else None
                corr_lift.append(hl if hl is not None else 0.0)
                ev_rows.append({"run": run, "outcome": outcome, "kind": kind, "chunk": d.get("after_chunk"),
                                "rel_to_close": None, "executed": True,
                                "h_state_mm": r1(None if not sp else H(sp) - h_close),
                                "h_actual_after_mm": r1(None if not ma else H(ma) - h_close),
                                "j1_measured": r1(None if not sp else sp[1]),
                                "j1_actual_after": r1(None if not ma else ma[1]),
                                "notes": f"人の補正 ({d.get('device') or 'キーボード'}) {d.get('action')}"
                                         + (f" grip {d.get('grip_before')}→{d.get('grip_after')}"),
                                })

        # ---- 閉じた瞬間の動き (2)
        def speed(t_a, t_b):
            qa, qb = q_at(act, t_a), q_at(act, t_b)
            return None if qa is None or qb is None else float(np.max(np.abs(qb - qa)) / (t_b - t_a))
        v_before = speed(t_close - 0.3, t_close)
        moved_during = (lambda qa, qb: None if qa is None or qb is None else float(np.max(np.abs(qb - qa))))(
            q_at(act, t_close), q_at(act, t_close + 0.8))
        moved_during_j = (lambda qa, qb: None if qa is None or qb is None else int(np.argmax(np.abs(qb - qa))))(
            q_at(act, t_close), q_at(act, t_close + 0.8))
        moved_2s = (lambda qa, qb: None if qa is None or qb is None else float(np.max(np.abs(qb - qa))))(
            q_at(act, t_close), q_at(act, t_close + 2.0))
        last_send = [s for s in sends if s[0] <= t_close]
        ls_t, ls_q = (last_send[-1][0], last_send[-1][1]) if last_send else (None, None)
        gap_close = None if ls_q is None else np.abs(ls_q - q_close)
        settled_q = q_at(act, t_close + 2.0)
        gap_settled = None if ls_q is None or settled_q is None else float(np.max(np.abs(ls_q - settled_q)))
        h_still = H(settled_q) - h_close if settled_q is not None else None
        # 届いた角度で, 閉じた瞬間から先に何 mm 爪先が動いたか (前後の位置は j0 と高さで)
        # 送信無視 (LOG352 / DAY228): 閉じたあと最初の送信に, 腕が向かい始めたか
        s_after = [s for s in sends if s[0] > t_close]
        ignored = None; first_send_info = ""
        if s_after:
            ts, qt, stg = s_after[0]
            q0 = q_at(act, ts)
            q2 = q_at(act, ts + 2.0)
            if q0 is not None and q2 is not None:
                d0 = float(np.max(np.abs(qt - q0))); d2 = float(np.max(np.abs(qt - q2)))
                moved = float(np.max(np.abs(q2 - q0)))
                # 無視された = 目標が 2° 以上離れているのに, 2 秒たっても 0.5° 未満しか動かない
                ignored = bool(d0 >= 2.0 and moved < 0.5)
                first_send_info = (f"{stg} 送信 {ts - t_close:.1f}s 後: 目標まで {d0:.1f}°→2秒後 {d2:.1f}° "
                                   f"(動いた {moved:.1f}°)")
        # 補正の1刻み (パッド・キーボード) で届かなかったもの
        corr_unreached = 0
        for d in log:
            if d.get("type") in ("pad", "correction") and (d.get("t_unix") or d.get("t") or 0) > t_close:
                tt = d.get("t_unix") if d.get("type") == "pad" else d.get("t")
                if not isinstance(tt, (int, float)) or tt > t_end_win:
                    continue
                if d.get("sent") and (d.get("arrived") is False or d.get("reached") is False):
                    corr_unreached += 1

        sg = {rel: (src, g, g1) for rel, src, g, g1 in state_gaps}
        close_rows.append({
            "run": run, "outcome": outcome, "close_stage": close_stage, "close_chunk": close_chunk,
            "t_close_s": round(t_close - act[0, 0], 2), "n_close_cmds": len(closes),
            "h_close_unity_mm": r1(h_close), "j1_close": r1(q_close[1]), "j2_close": r1(q_close[2]),
            "last_send_before_close_s": None if ls_t is None else round(t_close - ls_t, 2),
            "gap_lastsend_vs_actual_at_close_deg": None if gap_close is None else r1(gap_close.max()),
            "gap_joint": None if gap_close is None else f"j{int(np.argmax(gap_close))}",
            "gap_lastsend_vs_actual_2s_after_deg": r1(gap_settled),
            "speed_0.3s_before_close_deg_s": r1(v_before),
            "moved_during_grip_wait_0.8s_deg": r1(moved_during),
            "moved_during_grip_wait_joint": None if moved_during_j is None else f"j{moved_during_j}",
            "moved_0_to_2s_deg": r1(moved_2s),
            "h_change_0_to_2s_mm": r1(h_still),
            "first_send_after_close": first_send_info, "first_send_ignored": ignored,
            "corr_steps_unreached_after_close": corr_unreached,
            "state_src_close_chunk": sg.get(0, (None,))[0], "state_gap_close_chunk_deg": r1(sg.get(0, (0, None))[1]),
            "state_gap_close_chunk_j1_deg": r1(sg.get(0, (0, None, None))[2]),
            "state_src_next_chunk": sg.get(1, (None,))[0], "state_gap_next_chunk_deg": r1(sg.get(1, (0, None))[1]),
            "state_gap_next_chunk_j1_deg": r1(sg.get(1, (0, None, None))[2]),
            "state_src_prev_chunk": sg.get(-1, (None,))[0], "state_gap_prev_chunk_deg": r1(sg.get(-1, (0, None))[1]),
        })

        # ---- 持ち上げの区分 (1)
        raw = max(model_lift_raw) if model_lift_raw else None
        clip = max(model_lift_clip) if model_lift_clip else None
        sentm = max(model_lift_sent) if model_lift_sent else None
        actm = max(model_lift_act) if model_lift_act else None
        if raw is None:
            lift_cls = "推論なし (閉じてすぐ人の補正)"
        elif raw < LIFT_OUT_MM:
            lift_cls = "出していない"
        elif clip is not None and clip < LIFT_OUT_MM:
            lift_cls = "制限で削られた"
        elif sentm is None:
            lift_cls = "出したが送らなかった (人が止めた)"
        elif actm is not None and actm < LIFT_OUT_MM and sentm >= LIFT_OUT_MM:
            lift_cls = "送ったのに届かなかった"
        else:
            lift_cls = "出して届いた"
        sum_rows.append({
            "run": run, "outcome": outcome, "close_chunk": close_chunk,
            "lift_total_mm": r1(lift_total),
            "lift_by_model_actual_mm": r1(actm), "lift_by_correction_mm": r1(max(corr_lift) if corr_lift else None),
            "n_model_chunks_before_correction": len(model_lift_raw),
            "model_raw_max_mm": r1(raw), "model_clipped_max_mm": r1(clip), "model_sent_max_mm": r1(sentm),
            "model_actual_max_mm": r1(actm), "lift_class": lift_cls,
            "clamp_notes": " | ".join(f"c{c}: {'/'.join(n)}" for c, n in note_clamps),
            "b_category": btab.get(run, {}).get("category"),
            "b_close_tip_block_xz_mm": btab.get(run, {}).get("close_tip_block_xz_mm"),
            "b_lift_rise_mm": btab.get(run, {}).get("lift_rise_mm"),
            "judge_pick": tj.get("judgments", {}).get("pick", {}).get("human_verdict"),
        })

    def wcsv(name, rows):
        keys = []
        for r in rows:
            for k in r:
                if k not in keys:
                    keys.append(k)
        with open(OUT / name, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=keys); w.writeheader(); w.writerows(rows)
    wcsv("e_chunks.csv", ev_rows)
    wcsv("e_close.csv", close_rows)
    wcsv("e_summary.csv", sum_rows)
    print("fit", fit["holdout_abs_err_median_mm"], fit["holdout_abs_err_max_mm"])


if __name__ == "__main__":
    main()
