"""依頼110 E-4: 閉じる命令の直前・直後 (+1.0 秒)・π0.5 の持ち上げ1かたまり目を送り終えた直後・
爪先が閉じたときより 40mm 上がった (近似式. 送り手は π0.5 でも人の補正でもよい) 0.5 秒後 の上カメラのコマを並べる.

ブロックのあたり (直前のコマの赤の重心) を 420×315 px で切り出し, 1.25 倍にして並べる.
出力: e_grasp_failures.jpg (失敗12件) / e_grasp_successes.jpg (成功5件) / e_frames.csv (使ったコマ)
"""
import csv
import json
from pathlib import Path

import cv2
import numpy as np

ROOT = Path("/Volumes/KIOXIA/r2r2r-data")
RUNS = ROOT / "runs" / "pi05_policy"
TW = ROOT / "lab" / "twin" / "20261107_失敗の再生と置き場所"
OUT = TW / "E_つかみと持ち上げ"
CW, CH, SC = 420, 315, 1.25
FALLBACK = (1140, 690)        # 赤が見つからないときの中心 (17件の直前のコマの赤の位置のおおよそ)


def red_center(img):
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    m = cv2.inRange(hsv, (0, 120, 70), (10, 255, 255)) | cv2.inRange(hsv, (170, 120, 70), (180, 255, 255))
    m[:, :700] = 0; m[:, 1400:] = 0; m[:350] = 0
    n, lab, st, cen = cv2.connectedComponentsWithStats(m)
    best = None
    for k in range(1, n):
        x, y, w, h, a = st[k]
        if a < 300 or max(w, h) > 3 * min(w, h):      # 細長い赤 (配線) は外す
            continue
        if best is None or a > best[0]:
            best = (a, cen[k])
    return None if best is None else (int(best[1][0]), int(best[1][1]))


def grab(cap, idx):
    cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
    ok, f = cap.read()
    return f if ok else None


def crop(f, c):
    x0 = int(np.clip(c[0] - CW // 2, 0, f.shape[1] - CW)); y0 = int(np.clip(c[1] - CH // 2, 0, f.shape[0] - CH))
    im = f[y0:y0 + CH, x0:x0 + CW]
    return cv2.resize(im, (int(CW * SC), int(CH * SC)), interpolation=cv2.INTER_CUBIC)


def label(img, text, color=(255, 255, 255)):
    cv2.rectangle(img, (0, 0), (img.shape[1], 30), (0, 0, 0), -1)
    cv2.putText(img, text, (8, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.62, color, 2, cv2.LINE_AA)
    return img


def height_fn():
    c = np.array(json.load(open(OUT / "e_fk_fit.json"))["coef"])

    def H(q):
        j1, j2, j3, j4 = [np.radians(q[k]) for k in (1, 2, 3, 4)]
        b = j1 + j2
        f = np.array([1, 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)])
        return float(f @ c)
    return H


def t_rise(rec, tc, H, rise=40.0):
    """閉じたあと, 爪先が閉じたときより rise mm 上がった最初の時刻 (届いた角度. 次に開くまで)."""
    act, t_open = [], np.inf
    for r in csv.DictReader(open(rec / "states.csv")):
        if r["event_type"] == "gripper_cmd" and r["gripper_state"] == "0" and float(r["ts"]) > tc:
            t_open = min(t_open, float(r["ts"]))
        if r["event_type"] == "actual" and r["j0"] and r["stage"] != "home":     # HOME に戻る動きは数えない
            act.append((float(r["ts"]), [float(r[f"j{k}"]) for k in range(6)]))
    before = [q for t, q in act if t <= tc]
    if not before:
        return None
    h0 = H(before[-1])
    for t, q in act:
        if tc < t < t_open and H(q) - h0 >= rise:
            return t
    return None


def main():
    H = height_fn()
    close = {r["run"]: r for r in csv.DictReader(open(OUT / "e_close.csv"))}
    sel = json.load(open(TW / "B" / "selection.json"))
    rows_out = []
    for group, runs in (("failures", sel["failures"]), ("successes", sel["successes"])):
        tiles = []
        for run in runs:
            W, Hh = int(CW * SC), int(CH * SC)
            blank = np.full((Hh, W, 3), 40, np.uint8)
            rec = RUNS / run / "recording"
            fl = [(int(r["frame_idx"]), float(r["capture_ts"])) for r in csv.DictReader(open(rec / "top_frame_log.csv"))]
            gr = [float(r["ts"]) for r in csv.DictReader(open(rec / "states.csv"))
                  if r["event_type"] == "gripper_cmd" and r["gripper_state"] == "1"]
            name = run.replace("pi05_", "")
            if not gr:
                tiles.append(np.hstack([label(blank.copy(), f"{name} no close cmd")] + [blank] * 3))
                rows_out.append({"run": run, "note": "爪を閉じる命令の記録なし"})
                continue
            tc = gr[0]
            ts = np.array([t for _, t in fl])
            i_before = int(np.searchsorted(ts, tc, side="right") - 1)
            i_after = int(np.searchsorted(ts, tc + 1.0))
            # π0.5 の持ち上げ1かたまり目 (閉じたかたまり+1) を送り終えた直後
            cc = int(close[run]["close_chunk"])
            recs = {d["chunk"]: d for d in map(json.loads, open(RUNS / run / "log.jsonl"))
                    if d.get("type") is None and "chunk" in d}
            nx = recs.get(cc + 1)
            i_lift = None
            if nx and nx.get("executed") and nx.get("t_exec_end"):
                i_lift = int(np.searchsorted(ts, nx["t_exec_end"] + 0.3))
            cap = cv2.VideoCapture(str(rec / "top.mp4"))
            fb = grab(cap, fl[i_before][0]); fa = grab(cap, fl[min(i_after, len(fl) - 1)][0])
            c = red_center(fb) or red_center(fa) or FALLBACK
            t_b = ts[i_before] - tc; t_a = ts[min(i_after, len(fl) - 1)] - tc
            col = (90, 90, 255) if group == "failures" else (120, 255, 120)
            tb = label(crop(fb, c), f"{name} before {t_b:+.2f}s", col)
            ta = label(crop(fa, c), f"after {t_a:+.2f}s")
            if i_lift is not None and i_lift < len(fl):
                fl_img = grab(cap, fl[i_lift][0]); t_l = ts[i_lift] - tc
                tl = label(crop(fl_img, c), f"lift chunk {cc + 1} end {t_l:+.1f}s")
            else:
                tl = label(blank.copy(), f"chunk {cc + 1} not sent (human)")
                t_l = None
            tr = t_rise(rec, tc, H)
            i_rise = None if tr is None else int(np.searchsorted(ts, tr + 0.5))
            if i_rise is not None and i_rise < len(fl):
                t_r = ts[i_rise] - tc
                tr_img = label(crop(grab(cap, fl[i_rise][0]), c), f"tip +40mm {t_r:+.1f}s")
            else:
                t_r = None
                tr_img = label(blank.copy(), "tip never +40mm")
            tiles.append(np.hstack([tb, ta, tl, tr_img]))
            rows_out.append({"run": run, "group": group, "frame_before": fl[i_before][0], "t_before_s": round(t_b, 2),
                             "frame_after": fl[min(i_after, len(fl) - 1)][0], "t_after_s": round(t_a, 2),
                             "frame_lift": None if i_lift is None else fl[i_lift][0],
                             "t_lift_s": None if t_l is None else round(t_l, 2),
                             "frame_rise40": None if t_r is None else fl[i_rise][0],
                             "t_rise40_s": None if t_r is None else round(t_r, 2),
                             "crop_center_px": f"{c[0]},{c[1]}"})
            cap.release()
        cv2.imwrite(str(OUT / f"e_grasp_{group}.jpg"), np.vstack(tiles), [cv2.IMWRITE_JPEG_QUALITY, 88])
    keys = sorted({k for r in rows_out for k in r}, key=lambda k: list(rows_out[1].keys()).index(k)
                  if k in rows_out[1] else 99)
    with open(OUT / "e_frames.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=keys); w.writeheader(); w.writerows(rows_out)


if __name__ == "__main__":
    main()
