"""依頼113 G-3: 学習データ 30 本 (と scripted_v1 20 本) の「ラベルが 1 になる直前・直後」を, 元の top.mp4 の
同じ時刻のコマで E と同じ測り方 (E_つかみと持ち上げ/e_grasp_metrics.py の measure・e_images.py の red_center) で測る.

測るコマ (g_timing.csv の frame_* 列. 時刻は top_frame_log.csv の capture_ts):
  pre1p5 / pre0p5 : 爪の命令の 1.5 秒前・0.5 秒前 (まだ奥へ進んでいたか)
  before          : 爪の命令の直前 (E の before と同じ)
  ep_label        : エピソードでラベルが 1 になったコマに使われた画に近いコマ
  after           : 爪の命令の 1.0 秒後 (E の after と同じ. 爪が閉じきった形)
出力: g_grasp_metrics.csv / g_grasp_main.jpg (30 本. 1 行 1 本. 左から 1.5 秒前・直前・ラベルのコマ・1.0 秒後)
      g_grasp_scripted.jpg (scripted_v1)
"""
import csv
import sys
from pathlib import Path

import cv2
import numpy as np

ROOT = Path("/Volumes/KIOXIA/r2r2r-data")
TW = ROOT / "lab" / "twin" / "20261107_失敗の再生と置き場所"
OUT = TW / "G_お手本の閉じる時刻"
sys.path.insert(0, str(TW / "E_つかみと持ち上げ"))
from e_grasp_metrics import measure  # noqa: E402
from e_images import FALLBACK, CW, CH, SC, crop, label, red_center  # noqa: E402

RED_LO1, RED_HI1, RED_LO2, RED_HI2 = (0, 120, 70), (10, 255, 255), (170, 120, 70), (180, 255, 255)   # E と同じ
BLUE_LO, BLUE_HI = (100, 150, 80), (130, 255, 255)                                               # E と同じ
HALF = 350            # 爪の向きを測るときの切り出し (胴の青い印まで入るように E の 260 より広い)


def _red_mask(hsv, s_min):
    return cv2.inRange(hsv, (0, s_min, 70), (10, 255, 255)) | cv2.inRange(hsv, (170, s_min, 70), (180, 255, 255))


def red_center_wide(img, s_min=120):
    """E の red_center と同じ赤の決まり (面積 300 以上・細長くない) で, 探す範囲を画の全体に広げたもの.
    10/15〜10/19 の記録はブロックが画の左寄り (x 450〜800) にあり, E の範囲 (x 700〜1400) から外れるため.
    画の全体では赤い配線も拾うので, 爪先の青い印 2 つが 200px 以内にある赤の塊だけを候補にする."""
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    m = _red_mask(hsv, s_min)
    m[:, :40] = 0; m[:, -40:] = 0; m[:40] = 0; m[-40:] = 0
    n, lab, st, cen = cv2.connectedComponentsWithStats(m)
    nb, _, stb, cenb = cv2.connectedComponentsWithStats(cv2.inRange(hsv, BLUE_LO, BLUE_HI))
    blues = [cenb[k] for k in range(1, nb) if stb[k][4] >= 60]
    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
        d = sorted(np.hypot(*(b - cen[k])) for b in blues)
        if len(d) < 2 or d[1] > 200:
            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 _axis_once(roi, s_min):
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
    n, lab, st, cen = cv2.connectedComponentsWithStats(_red_mask(hsv, s_min))
    rb = [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])]
    nb, _, stb, cenb = cv2.connectedComponentsWithStats(cv2.inRange(hsv, BLUE_LO, BLUE_HI))
    bb = [k for k in range(1, nb) if stb[k][4] >= 60]
    if not rb or len(bb) < 2:
        return None
    r = max(rb, key=lambda k: st[k][4])
    bb = sorted(bb, key=lambda k: np.hypot(*(cenb[k] - cen[r])))
    c1, c2 = cenb[bb[0]], cenb[bb[1]]
    mid = (c1 + c2) / 2
    tip_dir = (c2 - c1) / max(1e-6, np.linalg.norm(c2 - c1))
    body = None
    for k in bb[2:]:
        v = cenb[k] - mid; d = np.linalg.norm(v)
        if 120 <= d <= 400 and abs(np.dot(v / d, tip_dir)) < 0.5:
            body = cenb[k]; break
    if body is not None:
        u = (mid - body) / np.linalg.norm(mid - body); src = "body"
    else:
        # 胴の印が写らないとき: 爪先の印 2 つを結ぶ線に直角で, 画の下向き (どの記録も爪先は画の下寄りを向く)
        u = np.array([-tip_dir[1], tip_dir[0]])
        u = u if u[1] > 0 else -u
        src = "tips"
    ys, xs = np.nonzero(lab == r)
    proj = (np.stack([xs, ys], 1) - mid) @ u
    lat = float((cen[r] - mid) @ np.array([-u[1], u[0]]))
    return {"axis_protrude_px": round(float(proj.max()), 1), "axis_lateral_px": round(lat, 1),
            "axis_marker_gap_px": round(float(np.linalg.norm(c1 - c2)), 1),
            "axis_deg": round(float(np.degrees(np.arctan2(u[0], u[1]))), 1), "axis_from": src,
            "axis_red_area": int(st[r][4]), "axis_s_min": s_min}


def measure_axis(f, c, s_min=120):
    """E の measure を爪の向きに沿って測るようにしたもの.
    爪の向き u = 胴の青い印 → 爪先の青い印 2 つの中点 (画の上で. 爪先へ向かう向き).
    胴の印が写らないときは, 爪先の印 2 つを結ぶ線に直角で画の下向き (axis_from=tips).
      axis_protrude_px : ブロックの赤の画素を u に投げた最大 − 爪先の印の中点 (E の protrude_px と同じ意味.
                         爪が画の真下を向いていれば E の値とほぼ同じになる)
      axis_lateral_px  : ブロックの重心の, u に直角な向きのずれ
      axis_marker_gap_px: 爪先の印の間 (E と同じ)
    赤の決まりは E と同じ (彩度 120 以上). ただしブロックが薄いピンクに写る回 (10/19 の一部) は彩度 120 では
    ブロックの一部しか拾えないので, 彩度 60 で拾った面積の半分未満なら彩度 60 の値を使う (axis_s_min=60)."""
    H_, W_ = f.shape[:2]
    x0, y0 = max(0, c[0] - HALF), max(0, c[1] - HALF)
    roi = f[y0:min(H_, c[1] + HALF), x0:min(W_, c[0] + HALF)]
    m = _axis_once(roi, s_min)
    if s_min == 120:
        m60 = _axis_once(roi, 60)
        if m60 and (m is None or m["axis_red_area"] < 0.5 * m60["axis_red_area"]):
            m = m60
    return m


TAGS = ("pre1p5", "pre0p5", "before", "ep_label", "after")
SHOW = (("pre1p5", "-1.5s"), ("before", "before"), ("ep_label", "label=1 frame"), ("after", "after +1.0s"))


def rec_of(run):
    sub = "pi05_policy" if run.startswith("pi05_") else "run_pick_place"
    return ROOT / "runs" / sub / run / "recording"


def main():
    rows = list(csv.DictReader(open(OUT / "g_timing.csv")))
    out, tiles = [], {"main": [], "scripted": []}
    W, Hh = int(CW * SC), int(CH * SC)
    for r in rows:
        if r["dataset"] == "teach_v1":            # trim2 と同じ記録・同じコマ (g_timing で確かめた) なので測らない
            continue
        cap = cv2.VideoCapture(str(rec_of(r["run"]) / "top.mp4"))
        fr = {}
        for tag in TAGS:
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(r[f"frame_{tag}"])); ok, f = cap.read()
            fr[tag] = f if ok else None
        ce = red_center(fr["before"]) or red_center(fr["after"])          # E と同じ探し方
        cw = None if ce else (red_center_wide(fr["before"]) or red_center_wide(fr["after"]))
        s_min = 120
        if not ce and not cw:      # ブロックが薄いピンクに写る回 (10/19 の一部). 彩度の下限だけ 60 に下げる
            s_min = 60
            cw = red_center_wide(fr["before"], 60) or red_center_wide(fr["after"], 60)
        c = ce or cw or FALLBACK
        o = {"dataset": r["dataset"], "run": r["run"], "kind": r["kind"], "crop_center_px": f"{c[0]},{c[1]}",
             "center_by": "E" if ce else ("wide" if cw and s_min == 120 else ("wide_s60" if cw else "none"))}
        for tag in TAGS:
            m = measure(fr[tag], c) if fr[tag] is not None else None
            for k, v in (m or {}).items():
                o[f"{tag}_{k}"] = v
            m = measure_axis(fr[tag], c, s_min) if fr[tag] is not None else None
            for k, v in (m or {}).items():
                o[f"{tag}_{k}"] = v
        out.append(o)
        name = r["run"].replace("run_0001_", "").replace("pi05_", "p")
        col = (255, 220, 120) if r["kind"] == "corr" else (120, 255, 120)
        row_t = []
        for j, (tag, txt) in enumerate(SHOW):
            img = crop(fr[tag], c) if fr[tag] is not None else np.full((Hh, W, 3), 40, np.uint8)
            p = o.get(f"{tag}_axis_protrude_px"); g = o.get(f"{tag}_axis_marker_gap_px")
            extra = "" if p is None else f" pr {p:+.0f} gap {g:.0f}"
            row_t.append(label(img, (f"{name} " if j == 0 else "") + txt + extra, col if j == 0 else (255, 255, 255)))
        tiles["main" if r["dataset"] == "corr_v5_main" else "scripted"].append(np.hstack(row_t))
        cap.release()
    # E の失敗・成功 16 件も同じ「爪の向きに沿う」測り方で測る (E の e_frames.csv の before / after のコマ)
    E = TW / "E_つかみと持ち上げ"
    grp = {r["run"]: r for r in csv.DictReader(open(E / "e_grasp_metrics.csv"))}
    for r in csv.DictReader(open(E / "e_frames.csv")):
        if not r.get("frame_after"):
            continue
        cap = cv2.VideoCapture(str(ROOT / "runs" / "pi05_policy" / r["run"] / "recording" / "top.mp4"))
        c = tuple(int(v) for v in r["crop_center_px"].split(","))
        o = {"dataset": "E_" + r["group"], "run": r["run"], "kind": "E", "crop_center_px": r["crop_center_px"],
             "center_by": "E"}
        for tag in ("before", "after"):
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(r[f"frame_{tag}"])); ok, f = cap.read()
            o[f"{tag}_protrude_px"] = grp[r["run"]].get(f"{tag}_protrude_px")
            o[f"{tag}_marker_gap_px"] = grp[r["run"]].get(f"{tag}_marker_gap_px")
            m = measure_axis(f, c) if ok else None
            for k, v in (m or {}).items():
                o[f"{tag}_{k}"] = v
        cap.release()
        out.append(o)
    for k, v in tiles.items():
        cv2.imwrite(str(OUT / f"g_grasp_{k}.jpg"), np.vstack(v), [cv2.IMWRITE_JPEG_QUALITY, 85])
    keys = []
    for o in out:
        keys += [k for k in o if k not in keys]
    with open(OUT / "g_grasp_metrics.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=keys); w.writeheader(); w.writerows(out)
    print(len(out), "rows")


if __name__ == "__main__":
    main()