"""依頼110 E-4 補足: 閉じたあと, ブロックが爪先から離れた時刻を上カメラで追う.

0.5 秒おきのコマで, 赤いブロックの重心と, 赤にいちばん近い青い印2つ (間が 60〜150px のもの = 爪先の印) の
中点の距離 sep を測る. 閉じた直後 (+1.0〜2.0 秒) の sep の中央値を基準にし, 基準より 40px 以上離れたコマが
2つ続いた最初の時刻を「離れた」とする. HOME に戻る区間と, 次に爪を開く命令のあとは見ない.
真上からの画なので, ブロックが爪といっしょに持ち上がったか (高さ) は分からない. 分かるのは横に離れたことだけ.
出力: e_track.csv (試行ごと) / e_track_series.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"
OUT = ROOT / "lab" / "twin" / "20261107_失敗の再生と置き場所" / "E_つかみと持ち上げ"
SEP_JUMP_PX = 40


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 find(f):
    hsv = cv2.cvtColor(f, cv2.COLOR_BGR2HSV)
    red = cv2.inRange(hsv, (0, 120, 70), (10, 255, 255)) | cv2.inRange(hsv, (170, 120, 70), (180, 255, 255))
    blue = cv2.inRange(hsv, (100, 150, 80), (130, 255, 255))
    n, _, st, cen = cv2.connectedComponentsWithStats(red)
    rs = [(st[k][4], cen[k]) for k in range(1, n)
          if st[k][4] >= 300 and max(st[k][2], st[k][3]) < 3 * min(st[k][2], st[k][3])]
    n, _, st, cen = cv2.connectedComponentsWithStats(blue)
    bs = [cen[k] for k in range(1, n) if st[k][4] >= 60]
    if not rs or len(bs) < 2:
        return None
    rc = max(rs, key=lambda x: x[0])[1]
    best = None
    for i in range(len(bs)):
        for j in range(i + 1, len(bs)):
            g = np.hypot(*(bs[i] - bs[j]))
            if 60 <= g <= 150:
                mid = (bs[i] + bs[j]) / 2
                d = np.hypot(*(rc - mid))
                if best is None or d < best[0]:
                    best = (d, g)
    return None if best is None else (best[0], best[1], rc)


def main():
    H = height_fn()
    sel = json.load(open(OUT.parent / "B" / "selection.json"))
    res, series = [], []
    for run in sel["failures"] + sel["successes"]:
        rec = RUNS / run / "recording"
        st = list(csv.DictReader(open(rec / "states.csv")))
        gr = [float(r["ts"]) for r in st if r["event_type"] == "gripper_cmd" and r["gripper_state"] == "1"]
        if not gr:
            res.append({"run": run, "note": "閉じる命令の記録なし"}); continue
        tc = gr[0]
        t_open = min([float(r["ts"]) for r in st if r["event_type"] == "gripper_cmd"
                      and r["gripper_state"] == "0" and float(r["ts"]) > tc] or [np.inf])
        t_home = min([float(r["ts"]) for r in st if r["event_type"] == "stage" and r["stage"] == "home"
                      and float(r["ts"]) > tc] or [np.inf])
        act = [(float(r["ts"]), [float(r[f"j{k}"]) for k in range(6)]) for r in st
               if r["event_type"] == "actual" and r["j0"]]
        at = np.array([a[0] for a in act])

        def q_at(t):
            k = max(0, int(np.searchsorted(at, t, side="right") - 1)); return act[k][1]
        h0, j00 = H(q_at(tc)), q_at(tc)[0]
        fl = [(int(r["frame_idx"]), float(r["capture_ts"])) for r in csv.DictReader(open(rec / "top_frame_log.csv"))]
        ts = np.array([t for _, t in fl])
        cap = cv2.VideoCapture(str(rec / "top.mp4"))
        t_end = min(t_open, t_home, ts[-1])
        base, lost, run_hits = [], None, 0
        for t in np.arange(tc + 1.0, t_end, 0.5):
            i = int(np.searchsorted(ts, t))
            if i >= len(fl):
                break
            cap.set(cv2.CAP_PROP_POS_FRAMES, fl[i][0]); ok, f = cap.read()
            m = find(f) if ok else None
            q = q_at(t)
            row = {"run": run, "t_s": round(t - tc, 1), "h_rise_mm": round(H(q) - h0, 1),
                   "j0_turn_deg": round(q[0] - j00, 1),
                   "sep_px": None if m is None else round(m[0], 1), "marker_gap_px": None if m is None else round(m[1], 1)}
            series.append(row)
            if m is None:
                continue
            if t - tc <= 2.0:
                base.append(m[0]); continue
            b = float(np.median(base)) if base else None
            if b is not None and lost is None:
                run_hits = run_hits + 1 if m[0] > b + SEP_JUMP_PX else 0
                if run_hits >= 2:
                    lost = row
        res.append({"run": run, "sep_base_px": None if not base else round(float(np.median(base)), 1),
                    "lost_t_s": None if lost is None else lost["t_s"],
                    "lost_h_rise_mm": None if lost is None else lost["h_rise_mm"],
                    "lost_j0_turn_deg": None if lost is None else lost["j0_turn_deg"],
                    "lost_sep_px": None if lost is None else lost["sep_px"],
                    "max_h_rise_until_lost_or_end_mm": max([s["h_rise_mm"] for s in series if s["run"] == run
                                                            and (lost is None or s["t_s"] <= lost["t_s"])] or [None]),
                    "watched_until_s": round(t_end - tc, 1),
                    "end_reason": "開く命令" if t_end == t_open else ("HOME" if t_end == t_home else "記録の終わり")})
        cap.release()
    for name, rows in (("e_track.csv", res), ("e_track_series.csv", series)):
        keys = []
        for r in rows:
            keys += [k for k in r if k not in keys]
        with open(OUT / name, "w", newline="") as fo:
            w = csv.DictWriter(fo, fieldnames=keys); w.writeheader(); w.writerows(rows)
    for r in res:
        print(r)


if __name__ == "__main__":
    main()