"""依頼113 G: 学習データ (teach_v1_trim2_corr_v5 の 30 本) で, 爪のラベル (state[:,6]) が 0→1 になる時刻と
アームの動き・元の記録の爪の命令 (gripper_cmd)・アームが止まった時刻を並べる. 比べるために teach_v1 (詰める前)
と scripted_v1 (20 本) も同じ形で出す. 読むだけ (データセット・runs には書かない).

時刻の戻し方 (エピソードの各コマが元の記録のどの時刻か):
  教えた動き (run_pick_place) : build_demo_episodes と同じ t0 = max(上カメラの最初, 実測の最初) に t_orig (trim2) / t (teach_v1・scripted_v1) を足す
  補正入り (pi05_policy)       : データと対で残したコード (corr_v5/code/) の select_steps を同じ引数 (speed20・trim_hold) で
                                 もう一度計算し (ファイルは書かない), 残ったコマの元の時刻を取る
  どちらも, 戻した時刻で作り直した state がエピソードの state と一致することを確かめる (state_check_max_deg)

出力: g_timing.csv (1 本 1 行) / g_steps.csv (ラベルの前後 ±15 コマのエピソードの中身) / g_fit.json (爪先の近似式)
"""
import csv
import json
import sys
from pathlib import Path

import numpy as np

ROOT = Path("/Volumes/KIOXIA/r2r2r-data")
DS = ROOT / "datasets"
TW = ROOT / "lab" / "twin" / "20261107_失敗の再生と置き場所"
OUT = TW / "G_お手本の閉じる時刻"
MAIN = "mecharm_redblock_box_teach_v1_trim2_corr_v5"
CODE = DS / "mecharm_redblock_box_corr_v5" / "code"
sys.path.insert(0, str(CODE))
import build_demo_episodes as bde  # noqa: E402
import build_correction_episodes as bce  # noqa: E402
from rewind_correct import excluded_windows  # noqa: E402

WIN = 1.5            # ラベルの前後で見る秒
STOP_TOL_DEG = 0.5   # 閉じる命令のときの姿勢からこの中にずっと入っていれば「止まっていた」
FPS = 10


# ---------------------------------------------------------------- 爪先の近似式 (高さは E の式. 横の伸びはここで当てはめ)
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_tip():
    ch = np.array(json.load(open(TW / "E_つかみと持ち上げ" / "e_fk_fit.json"))["coef"])
    B = TW / "B"
    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(1000 * np.hypot(float(r["tip_x"]), float(r["tip_z"])))   # 台の真上からの横の距離 (Unity)
    X, y = np.array(X), np.array(y)
    idx = np.random.default_rng(113).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 = {"reach": {"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),
                      "coef": [float(v) for v in c],
                      "note": "B の Unity の爪先 (tip_x, tip_z) の台からの横の距離を E と同じ j1〜j4 の三角関数で当てはめた. "
                              "台の中心は Unity の原点とした. 差 (どれだけ伸びたか) を見るためだけに使う"},
            "height": "E_つかみと持ち上げ/e_fk_fit.json の coef をそのまま使う"}
    json.dump(info, open(OUT / "g_fit.json", "w"), ensure_ascii=False, indent=1)
    return (lambda q: float((_feats(q) @ ch)[0])), (lambda q: float((_feats(q) @ c)[0]))


H, R = fit_tip()


# ---------------------------------------------------------------- 元の時刻を戻す
def demo_grid(rec):
    """build_demo_episodes.build_one と同じ前処理. 戻り値: 10Hz の時刻 t, 関節 q, ラベル g, 実測, 爪の命令."""
    at, aq, gt, gs = bde.load_states(rec / "states.csv")
    o = np.argsort(at); at, aq = at[o], aq[o]
    at, aq = bde.drop_spikes(at, aq)
    fidx, fts = bde.load_frame_log(rec / "top_frame_log.csv")
    t0, t1 = max(fts[0], at[0]), min(fts[-1], at[-1])
    t = np.arange(t0, t1, 1.0 / FPS)
    q = np.stack([np.interp(t, at, aq[:, k]) for k in range(6)], axis=1)
    return t, q, at, aq, gt, gs


def label_of(t, gt, gs):
    g = np.zeros(len(t))
    if len(gt):
        o = np.argsort(gt); gt, gs = gt[o], gs[o]
        for i, ti in enumerate(t):
            k = np.searchsorted(gt, ti, side="right") - 1
            g[i] = gs[k] if k >= 0 else 0.0
    return g


def corr_grid(run_dir):
    """build_correction_episodes.build_one と同じ計算 (corr_v5 の引数). 戻り値は demo_grid と同じ + 選んだ添字と状態."""
    rec = run_dir / "recording"
    at, aq, gt, gs = bde.load_states(rec / "states.csv")
    o = np.argsort(at); at, aq = at[o], aq[o]
    at, aq = bde.drop_spikes(at, aq)
    fidx, fts = bde.load_frame_log(rec / "top_frame_log.csv")
    t0, t1 = max(fts[0], at[0]), min(fts[-1], at[-1])
    t = np.arange(t0, t1, 1.0 / FPS)
    q, _ = bce.fill_motion_gaps(t, at, aq, mode="speed20")
    g = label_of(t, gt, gs)
    log = bce.load_log(run_dir / "log.jsonl")
    mode_ts, modes = bce.load_stage_times(rec / "states.csv")
    sel = bce.select_steps(t, q, g, mode_ts, modes, excluded_windows(log), trim_mid_idle=True, trim_hold=True)
    # 詰める前 (model / correct だけを残し, 止まりを詰める前) の添字
    from rewind_correct import step_keep_mask
    keep0, m_all = step_keep_mask(t, mode_ts, modes, excluded_windows(log))
    return t, q, at, aq, gt, gs, sel, np.flatnonzero(keep0), m_all


# ---------------------------------------------------------------- 1 本を調べる
def interp_q(t, q, x):
    return np.array([np.interp(x, t, q[:, k]) for k in range(6)])


def analyze(name, ep_dir, kind, ds_name):
    z = np.load(ep_dir / "episode.npz")
    meta = json.loads((ep_dir / "meta.json").read_text())
    st, ac = z["state"], z["actions"]
    tor = z["t_orig"] if "t_orig" in z.files else z["t"]
    if kind == "corr":
        run_dir = ROOT / "runs" / "pi05_policy" / name
        t, q, at, aq, gt, gs, sel, idx0, m_all = corr_grid(run_dir)
        absT = t[sel["idx"]]
        modes = meta.get("step_modes") or sel["modes"]
        pre_idx = idx0                                    # 詰める前に残っていたコマ (model / correct)
    else:
        run_dir = ROOT / "runs" / "run_pick_place" / name
        t, q, at, aq, gt, gs = demo_grid(run_dir / "recording")
        absT = t[0] + tor
        modes = ["teach"] * len(st)
        pre_idx = np.arange(len(t))
    rec = run_dir / "recording"
    chk = float(np.max(np.abs(np.stack([interp_q(t, q, x) for x in absT]) - st[:, :6])))
    sw = np.flatnonzero((st[1:, 6] == 1) & (st[:-1, 6] == 0)) + 1
    row = {"dataset": ds_name, "run": name, "kind": kind, "steps": len(st), "state_check_max_deg": round(chk, 3),
           "n_close_labels": len(sw)}
    if not len(sw):
        row["note"] = "ラベルが 0→1 にならない"
        return row, []
    i = int(sw[0])
    tl = float(absT[i])
    closes = sorted(float(x) for x, s in zip(gt, gs) if s == 1)
    tc = max([x for x in closes if x <= tl + 1e-6], default=np.nan)
    # アームが止まった時刻: 閉じる命令のときの姿勢 (直前の実測) から STOP_TOL_DEG の中に, 命令までずっと入っていた最初の実測
    kb = np.flatnonzero(at <= tc)
    qc = aq[kb[-1]]
    s = kb[-1]
    while s - 1 >= 0 and np.max(np.abs(aq[s - 1] - qc)) <= STOP_TOL_DEG:
        s -= 1
    t_stop = float(at[s])
    gap_before_stop = float(at[s] - at[s - 1]) if s > 0 else np.nan     # 止まった時刻の不確かさ (実測の間)
    after = (at > tc) & (at <= tc + WIN)
    move_after = float(np.max(np.abs(aq[after] - qc))) if after.any() else np.nan
    qa, q0, qb = interp_q(t, q, tl - WIN), interp_q(t, q, tl), interp_q(t, q, tl + WIN)
    qh = interp_q(t, q, tl - 0.5)
    # 閉じる前 0.3 秒の速さ (E と同じ見方. 実測の差の最大 / 0.3 秒)
    q03 = interp_q(t, q, tc - 0.3)
    qtc = interp_q(t, q, tc)
    row.update({
        "label_step": i, "label_mode": modes[i], "label_mode_prev": modes[i - 1],
        "t_cmd": round(tc, 3), "t_label_minus_cmd_s": round(tl - tc, 3),
        "t_cmd_minus_stop_s": round(tc - t_stop, 2), "stop_gap_before_s": round(gap_before_stop, 2),
        "move_after_cmd_1p5s_deg": None if np.isnan(move_after) else round(move_after, 2),
        "speed_pre0p3_deg_s": round(float(np.max(np.abs(qtc - q03))) / 0.3, 2),
        **{f"d_j{k}_pre": round(float(q0[k] - qa[k]), 2) for k in (1, 2, 4)},
        **{f"d_j{k}_post": round(float(qb[k] - q0[k]), 2) for k in (1, 2, 4)},
        "d_reach_pre_mm": round(float(R(q0) - R(qa)), 1), "d_reach_last0p5_mm": round(float(R(q0) - R(qh)), 1),
        "d_reach_post_mm": round(float(R(qb) - R(q0)), 1),
        "d_h_pre_mm": round(float(H(q0) - H(qa)), 1), "d_h_last0p5_mm": round(float(H(q0) - H(qh)), 1),
        "d_h_post_mm": round(float(H(qb) - H(q0)), 1),
        "h_close_mm": round(float(H(q0)), 1), "reach_close_mm": round(float(R(q0)), 1),
    })
    # 詰める処理: ラベルの前後 1.5 秒に元の 10Hz で何コマあったか / エピソードに何コマ残ったか
    near_pre = (t[pre_idx] >= tl - WIN) & (t[pre_idx] <= tl + WIN)
    near_ep = (absT >= tl - WIN - 1e-6) & (absT <= tl + WIN + 1e-6)
    row["grid_steps_pm1p5"] = int(near_pre.sum())
    row["ep_steps_pm1p5"] = int(near_ep.sum())
    lo = max(0, i - 15)
    row["ep_15steps_before_span_s"] = round(float(tl - absT[lo]), 1)     # エピソードの 15 コマ前が何秒前か
    row["ep_15steps_d_reach_mm"] = round(float(R(st[i, :6]) - R(st[lo, :6])), 1)
    row["ep_15steps_d_h_mm"] = round(float(H(st[i, :6]) - H(st[lo, :6])), 1)
    # ラベルの位置がずれていないか: エピソードで最初に 1 になるコマの元の時刻が, 10Hz の並びで最初に 1 になる時刻と同じか
    g_full = label_of(t, gt, gs)
    first_full = t[np.flatnonzero((g_full[1:] == 1) & (g_full[:-1] == 0))[0] + 1] if (g_full == 1).any() else np.nan
    row["t_label_minus_grid_first1_s"] = round(tl - float(first_full), 2)
    # 行動: ラベルの前後のコマで「次の目標」がどちらへ向いているか
    steps = []
    for k in range(max(0, i - 15), min(len(st), i + 16)):
        steps.append({"dataset": ds_name, "run": name, "rel": k - i, "t_rel_s": round(float(absT[k] - tl), 2),
                      "mode": modes[k], "grip_state": int(st[k, 6]), "grip_action": int(ac[k, 6]),
                      "act_max_joint_deg": round(float(np.max(np.abs(ac[k, :6] - st[k, :6]))), 2),
                      "act_d_reach_mm": round(float(R(ac[k, :6]) - R(st[k, :6])), 2),
                      "act_d_h_mm": round(float(H(ac[k, :6]) - H(st[k, :6])), 2),
                      **{f"act_d_j{j}": round(float(ac[k, j] - st[k, j]), 2) for j in (1, 2, 4)},
                      "reach_mm": round(float(R(st[k, :6])), 1), "h_mm": round(float(H(st[k, :6])), 1)})
    by = {s["rel"]: s for s in steps}
    for r in (-3, -2, -1, 0, 1):
        if r in by:
            row[f"act{r:+d}_d_reach_mm"] = by[r]["act_d_reach_mm"]
            row[f"act{r:+d}_max_joint_deg"] = by[r]["act_max_joint_deg"]
    pre = [s for s in steps if -10 <= s["rel"] < 0]
    row["act_pre10_sum_d_reach_mm"] = round(sum(s["act_d_reach_mm"] for s in pre), 1)
    row["act_pre10_sum_d_h_mm"] = round(sum(s["act_d_h_mm"] for s in pre), 1)
    row["frame_ts_before"] = None
    # 上カメラ: 爪の命令の直前のコマ (E と同じ) / エピソードのラベル直前・直後のコマに近いコマ
    fl = [(int(r["frame_idx"]), float(r["capture_ts"])) for r in csv.DictReader(open(rec / "top_frame_log.csv"))]
    fts = np.array([x for _, x in fl])

    def frame_at_or_before(x):
        return fl[max(0, int(np.searchsorted(fts, x, side="right") - 1))]

    def frame_after(x):
        return fl[min(len(fl) - 1, int(np.searchsorted(fts, x)))]

    def nearest(x):
        k = int(np.argmin(np.abs(fts - x)))
        return fl[k]
    row.update({
        "frame_pre1p5": frame_at_or_before(tc - WIN)[0], "frame_pre0p5": frame_at_or_before(tc - 0.5)[0],
        "frame_before": frame_at_or_before(tc)[0], "frame_after": frame_after(tc + 1.0)[0],
        "frame_ep_prev": nearest(float(absT[i - 1]))[0], "frame_ep_label": nearest(tl)[0],
        "t_before_s": round(frame_at_or_before(tc)[1] - tc, 2), "t_after_s": round(frame_after(tc + 1.0)[1] - tc, 2),
    })
    del row["frame_ts_before"]
    # 補正入り: 閉じる前に人の補正があったか (ラベルより前の correct のコマ) と, その間の伸び
    if kind == "corr":
        ci = [k for k in range(i) if modes[k] == "correct"]
        row["correct_steps_before_label"] = len(ci)
        if ci:
            k0, k1 = ci[0], ci[-1]
            row["correct_last_to_label_s"] = round(float(tl - absT[k1]), 1)
            row["correct_d_reach_mm"] = round(float(R(st[min(k1 + 1, len(st) - 1), :6]) - R(st[k0, :6])), 1)
            row["correct_d_h_mm"] = round(float(H(st[min(k1 + 1, len(st) - 1), :6]) - H(st[k0, :6])), 1)
            dj = st[min(k1 + 1, len(st) - 1), :6] - st[k0, :6]
            row["correct_d_joints"] = " ".join(f"{v:+.1f}" for v in dj)
    return row, steps


def episodes(ds):
    return sorted(p for p in (DS / ds / "raw_episodes").iterdir() if (p / "episode.npz").exists())


def main():
    rows, steps = [], []
    for p in episodes(MAIN):
        kind = "corr" if p.name.startswith("pi05_") else "teach"
        r, s = analyze(p.name, p, kind, "corr_v5_main"); rows.append(r); steps += s
    for ds, tag in (("mecharm_redblock_box_teach_v1", "teach_v1"), ("mecharm_redblock_box_scripted_v1", "scripted_v1")):
        for p in episodes(ds):
            r, s = analyze(p.name, p, "teach", tag); rows.append(r); steps += s
    keys = []
    for r in rows:
        keys += [k for k in r if k not in keys]
    with open(OUT / "g_timing.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=keys); w.writeheader(); w.writerows(rows)
    with open(OUT / "g_steps.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(steps[0])); w.writeheader(); w.writerows(steps)
    print(f"rows={len(rows)} steps={len(steps)}  state_check_max_deg={max(r['state_check_max_deg'] for r in rows)}")


if __name__ == "__main__":
    main()
