#!/usr/bin/env python3
"""make_images.py — LOG479 (フィジカルAI研究室 Log#26) の記事の画像を作る. 元の記録は読むだけ.
書き出し先: 第1引数のフォルダ (既定は iCloud の フィジカルAI研究室を作る/LOG479_images/)
  LOG479_01_pos3_old_vs_new.jpg  位置3 の閉じて 1.0 秒後の上カメラ. 左: 元の記録 (120818・42.5px), 右: 教え直し (124117・2.7px)
  LOG479_03_pos3_align.jpg       位置3 の最初のコマにブロックの枠を重ねた図 (緑: 深かった本, 赤: 浅かった本)
  LOG479_04_pos4_descent.jpg     位置4 の下ろす区間の上カメラ 5 コマ × 3 段 (125825 差し替え版・193825 10/19・130542 教え直し)
  LOG479_02_measure_cable.jpg    measure の確認画像 (measure_20261109_123559.jpg) を 1280 幅に縮めたもの
"""
import bisect, csv, json, sys
from pathlib import Path
import cv2
import numpy as np

DATA = Path("/Volumes/KIOXIA/r2r2r-data")
OUT = Path(sys.argv[1]) if len(sys.argv) > 1 else Path.home() / "Library/Mobile Documents/com~apple~CloudDocs/デスクトップ/ADAPT DESIGN/じぶんの/PixelstoMotors/デジタルツイン学習R2R2R/フィジカルAI研究室を作る/LOG479_images"
if len(sys.argv) > 2:
    DATA = Path(sys.argv[2])
RUNS = DATA / "runs/run_pick_place"
OUT.mkdir(parents=True, exist_ok=True)


def frame_at(run, t):
    ts = [float(r["capture_ts"]) for r in csv.DictReader(open(RUNS / run / "recording/top_frame_log.csv"))]
    cap = cv2.VideoCapture(str(RUNS / run / "recording/top.mp4"))
    cap.set(cv2.CAP_PROP_POS_FRAMES, min(bisect.bisect_left(ts, t), len(ts) - 1))
    ok, f = cap.read()
    cap.release()
    return f


def stage_t(run, name):
    st = list(csv.DictReader(open(RUNS / run / "recording/states.csv")))
    return [float(r["ts"]) for r in st if r["event_type"] == "stage" and r["stage"] == name][0]


def label(img, text):
    cv2.rectangle(img, (0, 0), (img.shape[1], 34), (20, 20, 20), -1)
    cv2.putText(img, text, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (240, 240, 240), 2, cv2.LINE_AA)
    return img


# 1. 位置3: 元の記録と教え直しの, 閉じて 1.0 秒後 (grasp_depth.py が測るコマ)
sys.path.insert(0, str(Path.home() / "r2r2r-scripts/pi05"))
import grasp_depth as gd  # noqa: E402
tiles = []
for run, txt in [("run_0001_20261109_120818", "old teach  42.5px"), ("run_0001_20261109_124117", "re-taught  2.7px")]:
    m = gd.measure_run(RUNS / run)
    cap = cv2.VideoCapture(str(RUNS / run / "recording/top.mp4"))
    cap.set(cv2.CAP_PROP_POS_FRAMES, m["frame_after"])
    ok, f = cap.read()
    cap.release()
    cx, cy = m["center_px"]
    c = f[max(0, cy - 260):cy + 260, max(0, cx - 330):cx + 330].copy()
    tiles.append(label(cv2.resize(c, (660, 520)), txt))
cv2.imwrite(str(OUT / "LOG479_01_pos3_old_vs_new.jpg"), np.hstack(tiles), [cv2.IMWRITE_JPEG_QUALITY, 90])

# 2. measure の確認画像 (ケーブルを拾った回)
f = cv2.imread(str(DATA / "mecharm_calibration/teach_poses_20261109/measure_20261109_123559.jpg"))
cv2.imwrite(str(OUT / "LOG479_02_measure_cable.jpg"), cv2.resize(f, (1280, 720)), [cv2.IMWRITE_JPEG_QUALITY, 90])

# 3. 位置3 の置き場所 (最初のコマのブロックの枠)
runs = {"124117": ("deep 2.7", (0, 200, 0)), "124256": ("deep 4.5", (0, 200, 0)), "123848": ("deep 9.5", (0, 200, 0)),
        "123655": ("shallow 25.9", (0, 0, 255)), "123940": ("shallow 25.6", (0, 0, 255))}
base = None
for k, (r, (txt, col)) in enumerate(runs.items()):
    cap = cv2.VideoCapture(str(RUNS / f"run_0001_20261109_{r}/recording/top.mp4"))
    ok, f = cap.read()
    cap.release()
    if base is None:
        base = f.copy()
    hsv = cv2.cvtColor(f, cv2.COLOR_BGR2HSV)
    m = cv2.inRange(hsv, (0, 40, 80), (12, 255, 255)) | cv2.inRange(hsv, (158, 40, 80), (180, 255, 255))
    m[:500, :] = 0   # 根元とケーブルを除く
    m = cv2.morphologyEx(m, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8))
    cs, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    rect = cv2.minAreaRect(max(cs, key=cv2.contourArea))
    cv2.drawContours(base, [cv2.boxPoints(rect).astype(int)], 0, col, 2)
    cx, cy = map(int, rect[0])
    cv2.putText(base, f"{r} {txt}", (cx + 60, cy - 40 + 20 * k), cv2.FONT_HERSHEY_SIMPLEX, 0.6, col, 2)
crop = cv2.resize(base[480:900, 420:1000], None, fx=1.6, fy=1.6)
cv2.imwrite(str(OUT / "LOG479_03_pos3_align.jpg"), crop, [cv2.IMWRITE_JPEG_QUALITY, 90])

# 4. 位置4 の下ろす区間 (PICK の段階の 0.5 秒前から, 段階の終わりの 2 秒前まで)
rows = []
for run, txt in [("run_0001_20261109_125825", "125825 (13:00) "), ("run_0001_20261019_193825", "193825 (10/19) "),
                 ("run_0001_20261109_130542", "130542 (13:03) ")]:
    tp, tl = stage_t(run, "PICK"), stage_t(run, "APPROACH_PICK_LIFT")
    row = []
    for t in [tp - 0.5] + list(np.linspace(tp + 0.5, max(tp + 0.6, tl - 2.0), 4)):
        f = frame_at(run, t)
        c = cv2.resize(f[100:700, 300:1000], None, fx=0.5, fy=0.5)
        cv2.putText(c, f"{txt}{t - tp:+.1f}s", (5, 20), 0, 0.6, (0, 255, 255), 2)
        row.append(c)
    rows.append(np.hstack(row))
cv2.imwrite(str(OUT / "LOG479_04_pos4_descent.jpg"), np.vstack(rows), [cv2.IMWRITE_JPEG_QUALITY, 90])
print("ok", sorted(p.name for p in OUT.glob("LOG479_*")))