#!/usr/bin/env python3
"""LOG484 / 研究室 Log28: deep1 の実機推論で、アームがどこをつかみに行ったかを、お手本の位置ごとのつかむ姿勢と比べる.

- お手本: mecharm_redblock_box_deep_even10_v1 の 50 本. 位置は pi05/count_by_position.py の出力 (2026-11-10) のとおり.
  つかむ姿勢 = state の爪 (7 列目) が初めて 1 になったコマの j0〜j5.
- 実機推論: runs/pi05_policy/pi05_20261110_*/.
  爪を閉じた姿勢 = recording/states.csv で, 初めて gripper_cmd が 1 (閉じる) になった直前の actual の j0〜j5.
  AI が行った所 = 閉じる前にゲームパッドの補正があれば, その補正の最初の measured_before (人が直す前の姿勢).
  なければ, 爪を閉じた姿勢と同じ.
- 近い位置 = AI が行った所と, 位置ごとのつかむ姿勢の平均との距離 (j0〜j5 の差の二乗和の平方根) がいちばん小さい位置.
使い方: python3 grasp_pose_compare.py [--data-root /Volumes/KIOXIA/r2r2r-data]
"""
import argparse, csv, json, math
from pathlib import Path
import numpy as np

POS = {
 1: "094839 095308 095810 095948 100126 100307 100444 100716 101157 103049",
 2: "191257 191541 191635 102109 102248 102431 103804 103943 104110 104517",
 3: "192438 192539 123848 124032 124117 124207 124256 124438 124817 124920",
 4: "193404 193459 193617 193825 193932 130542 130723 130824 130929 131103",
 5: "pi05_20261031_175527 pi05_20261031_180435 pi05_20261031_184016 pi05_20261101_190347 pi05_20261106_183926 pi05_20261106_184650 194818 194913 195023 195120",
}
DAY = {"1910": "20261019", "1930": "20261019", "1940": "20261019", "1950": "20261019"}
# 2026-11-10 の実機推論. 位置はエイイチの答え (チャット) とスクリプトの出力から. 111034 は位置合わせの前の 1 回目
RUNS = [("111034", "1回目(位置なし)"), ("112346", "位置1"), ("112639", "位置1"), ("112924", "位置1"),
        ("113433", "位置2"), ("113600", "位置2"), ("113741", "位置2"), ("114106", "位置2")]


def ep_name(s):
    if s.startswith("pi05_"):
        return s
    day = "20261019" if s[:2] == "19" else "20261109"
    return f"run_0001_{day}_{s}"


def demo_grasp(root):
    out = {}
    for p, names in POS.items():
        poses = []
        for s in names.split():
            z = np.load(root / "datasets/mecharm_redblock_box_deep_even10_v1/raw_episodes" / ep_name(s) / "episode.npz")
            st = z["state"]
            i = int(np.argmax(st[:, 6] > 0.5))
            poses.append(st[i, :6])
        out[p] = np.array(poses)
    return out


def run_poses(rd):
    rows = list(csv.DictReader(open(rd / "recording/states.csv")))
    last = None
    close = None
    close_t = None
    for r in rows:
        if r["event_type"] == "actual" and r["j0"]:
            last = [float(r[f"j{k}"]) for k in range(6)]
        if r["event_type"] == "gripper_cmd" and r["gripper_state"] == "1" and close is None:
            close, close_t = last, float(r["ts"])
    pad = None
    for line in open(rd / "log.jsonl"):
        d = json.loads(line)
        if d.get("type") == "pad" and d.get("measured_before"):
            if close_t is None or d["t_unix"] < close_t:
                pad = d["measured_before"]
            break
    meta = json.load(open(rd / f"{rd.name}.json"))
    return close, pad, meta


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-root", default="/Volumes/KIOXIA/r2r2r-data")
    a = ap.parse_args()
    root = Path(a.data_root)
    g = demo_grasp(root)
    mean = {p: v.mean(0) for p, v in g.items()}
    print("お手本 50 本の, 位置ごとのつかむ姿勢 (平均 / 最小〜最大)")
    print("位置 | " + " | ".join(f"j{k}" for k in range(6)))
    for p, v in g.items():
        print(f"{p} | " + " | ".join(f"{v[:, k].mean():+.1f} ({v[:, k].min():+.1f}〜{v[:, k].max():+.1f})" for k in range(6)))
    print()
    print("位置どうしの, つかむ姿勢の平均の距離 (度)")
    for p in mean:
        print(f"{p} | " + " | ".join(f"{np.linalg.norm(mean[p] - mean[q]):.1f}" for q in mean))
    print()
    print("2026-11-10 の実機推論 (deep1)")
    print("run | 置いた所 | 結果 | 補正 | AI が行った所 j0〜j5 | 近い位置 (距離) | 位置ごとの距離 1〜5 | 爪を閉じた姿勢 j0〜j5")
    for stem, label in RUNS:
        rd = root / "runs/pi05_policy" / f"pi05_20261110_{stem}"
        close, pad, meta = run_poses(rd)
        went = pad if pad is not None else close
        corr = meta.get("corrections", 0)
        res = meta.get("outcome") or meta.get("result")
        if went is None:
            print(f"{stem} | {label} | {res} | {corr} | 爪を閉じる前に終えた | - | - | -")
            continue
        dist = {p: float(np.linalg.norm(np.array(went) - mean[p])) for p in mean}
        near = min(dist, key=dist.get)
        cl = " ".join(f"{x:+.1f}" for x in close) if close else "閉じていない"
        print(f"{stem} | {label} | {res} | {corr} | " + " ".join(f"{x:+.1f}" for x in went)
              + f" | {near} ({dist[near]:.1f}) | " + " ".join(f"{dist[p]:.1f}" for p in mean) + f" | {cl}"
              + (" | 補正の前" if pad is not None else ""))


if __name__ == "__main__":
    main()