# 位置D・位置E の実機の回 (LOG363) を, 同じ位置のお手本 (D 5本・E 4本) と比べる
#  1) つかんだ姿勢 / 運ぶ間に肩がいちばん上がった所 (j1 の最小) / 離した姿勢
#  2) ブロックを持って肩を上げる指令の, 目標と次のチャンクの実測の差 (+ は目標より低いまま止まった)
#  3) 爪先の高さ (FK 表 fk_scan_motion_based_v2.csv の ee_y を (j1,j2,j3,j4) の近い 40 点で局所線形に当てはめる, LOG361 と同じ)
#  実機は中継点ごとに「送った目標 (target)」と「着いたあとの実測 (actual_after)」が残っているので, 肩の高さと爪先の高さは両方を出す
import os, json
import numpy as np, pandas as pd
ROOT = os.environ.get("R2R2R_DATA", "/Volumes/KIOXIA/r2r2r-data"); os.chdir(ROOT)
RAW = "datasets/mecharm_redblock_box_teach_v1_trim2/raw_episodes"
DEMO = {"D": ["193404", "193459", "193617", "193825", "193932"],
        "E": ["194818", "194913", "195023", "195120"]}
RUN = {"D": "20261024_190239", "E": "20261024_191603"}
d = pd.read_csv("fk_tables/fk_scan_motion_based_v2.csv")
X = d[["j1", "j2", "j3", "j4"]].to_numpy(float); Y = d["ee_y"].to_numpy(float)
def h(q):
    q = np.asarray(q, float)[[1, 2, 3, 4]]
    dd = ((X - q) ** 2).sum(1); idx = np.argpartition(dd, 40)[:40]
    A = np.c_[np.ones(len(idx)), X[idx] - q]
    coef, *_ = np.linalg.lstsq(A, Y[idx], rcond=None)
    return coef[0] * 1000  # mm
def p(v, n=6):
    return " ".join(f"{x:+6.1f}" for x in v[:n])
for pos in ["D", "E"]:
    print(f"\n######## 位置{pos}")
    print("== お手本 (閉じた瞬間 / 持っている間の j1 最小 / 開いた瞬間, 角度は j0..j5)")
    for e in DEMO[pos]:
        S = np.load(f"{RAW}/run_0001_20261019_{e}/episode.npz")["state"]
        g = S[:, 6]; c = int(np.argmax(g > 0.5)); o = c + int(np.argmax(g[c:] < 0.5))
        held = S[c:o]
        k = int(np.argmin(held[:, 1]))
        hs = [h(S[t]) for t in range(c, o + 1, 10)]
        print(f"demo {e}  閉 [{p(S[c])}]  j1最小 {held[k,1]:+.1f} (j0 {held[k,0]:+.1f})  開 [{p(S[o])}]")
        print(f"            爪先の高さ mm  閉じた時 {h(S[c]):6.1f}  最高 {max(h(S[t]) for t in range(c, o+1)):6.1f}  1秒ごと " + " ".join(f"{v:5.0f}" for v in hs))
    L = [json.loads(l) for l in open(f"runs/pi05_policy/pi05_{RUN[pos]}/log.jsonl")]
    L = [x for x in L if "chunk" in x]
    print(f"== 実機 pi05_{RUN[pos]}")
    close = [x for x in L if x.get("grip_cmd") == 1]
    opn = [x for x in L if x.get("grip_cmd") == 0]
    print(f"  閉じる指令のチャンク {[x['chunk'] for x in close]}  開く指令のチャンク {[x['chunk'] for x in opn]}")
    if close:
        c = close[0]
        print(f"  閉じた所 (閉じる指令のチャンクの最後の目標) [{p(c['target'])}]")
    held = [x for x in L if x["state"][6] == 1]
    meas = np.array([x.get("measured_state", x["state"])[:6] for x in held])
    vias = [v for x in held for v in (x.get("vias") or [])]
    vt = np.array([v["target"] for v in vias], float)
    va = np.array([v["actual_after"] for v in vias], float)
    k = int(np.argmin(meas[:, 1]))
    print(f"  持っている間 チャンクの頭の実測 j1 最小 {meas[k,1]:+.1f} (チャンク {held[k]['chunk']}, j0 {meas[k,0]:+.1f})")
    kk = int(np.argmin(vt[:, 1])); ka = int(np.argmin(va[:, 1]))
    print(f"  持っている間 中継点の目標 j1 最小 {vt[kk,1]:+.1f} (j0 {vt[kk,0]:+.1f})  中継点の実測 j1 最小 {va[ka,1]:+.1f} (j0 {va[ka,0]:+.1f})")
    print(f"  中継点の目標と実測の j1 の差 (+ は低いまま) 平均 {np.mean(va[:,1]-vt[:,1]):+.2f}  最小 {np.min(va[:,1]-vt[:,1]):+.2f}  最大 {np.max(va[:,1]-vt[:,1]):+.2f}  ({len(va)} 点)")
    print("  中継点ごと (j0 目標 / j1 目標 → 実測 / 爪先の高さ mm 目標 → 実測)")
    for v in vias:
        t = v["target"]; a = v["actual_after"]
        print(f"   j0 {t[0]:+6.1f}  j1 {t[1]:+6.1f} → {a[1]:+6.1f}  高さ {h(t):5.0f} → {h(a):5.0f}")
    print(f"  爪先の高さの最高 (中継点の実測) {max(h(a) for a in va):.1f} mm  (中継点の目標) {max(h(t) for t in vt):.1f} mm")
    if opn:
        o = opn[0]
        print(f"  離した所 (開く指令のチャンクの最後の目標) [{p(o['target'])}]")
    print("  爪先の高さ mm (持っている間のチャンクの頭の実測) " + " ".join(f"{h(s):5.0f}" for s in meas) + f"  最高 {max(h(s) for s in meas):.1f}")
    print("  肩 j1: 目標 (チャンクの最後) と 次のチャンクの頭の実測 (+ は目標より低いまま止まった)")
    for a, b in zip(L, L[1:]):
        if not a.get("executed"):
            continue
        s = np.array(a.get("measured_state", a["state"])[:6]); t = np.array(a["target"])
        n = np.array(b.get("measured_state", b["state"])[:6])
        print(f"   {a['chunk']:02d} 爪 {int(a['state'][6])}  指令 {t[1]-s[1]:+6.1f}  目標 {t[1]:+6.1f}  実測 {n[1]:+6.1f}  ずれ {n[1]-t[1]:+5.2f}")