#!/usr/bin/env python3
"""make_images.py — LOG463 の記事に載せる画像を、上カメラの録画 (top.mp4) と横からのスマホ動画から作る (読むだけ).
使い方: python3 make_images.py   (B と SIDE のパスを環境に合わせる)
"""
import cv2, csv, os, subprocess
import numpy as np
from PIL import Image, ImageDraw, ImageFont
FP = "/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc"
F = lambda s: ImageFont.truetype(FP, s, index=0)
B = os.environ.get("RUNS", "/mnt/user-data/uploads/KIOXIA/r2r2r-data/runs/run_pick_place") + "/run_0001_20261109_"
SIDE = os.environ.get("SIDE", "/mnt/user-data/attachments/13cd2fb7-9204-537e-90d7-7ae29d9f967a")
OUT = "LOG463_images"


def times(s):
    st = list(csv.DictReader(open(B + s + "/recording/states.csv")))
    tp = [float(r["ts"]) for r in st if r["event_type"] == "stage" and r["stage"] == "PICK"][0]
    tc = [float(r["ts"]) for r in st if r["event_type"] == "gripper_cmd" and r["gripper_state"] == "1"][0]
    return tp, tc


def frame(s, t):
    fl = [(int(r["frame_idx"]), float(r["capture_ts"])) for r in csv.DictReader(open(B + s + "/recording/top_frame_log.csv"))]
    fi = min(fl, key=lambda x: abs(x[1] - t))[0]
    cap = cv2.VideoCapture(B + s + "/recording/top.mp4"); cap.set(cv2.CAP_PROP_POS_FRAMES, fi)
    ok, im = cap.read(); return im


def pil(im): return Image.fromarray(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))


# 04: 前の教えた記録と作り直した記録. 閉じて 1.0 秒後の上カメラ
def img04():
    pairs = [("093045", "前の教えた記録（2022-10-13 に手で教えた）", "閉じて1秒後のはみ出し 43.4px（浅い）", (200, 30, 30)),
             ("094839", "ゲームパッドで作り直した記録", "閉じて1秒後のはみ出し 16.2px（深い）", (30, 90, 200))]
    W = 2 * 660 + 60; out = Image.new("RGB", (W, 564 + 110), "white"); d = ImageDraw.Draw(out)
    for i, (s, t1, t2, col) in enumerate(pairs):
        tp, tc = times(s); im = frame(s, tc + 1.0)
        c = cv2.resize(im[330:800, 330:880], (660, 564)); x = 20 + i * 680
        out.paste(pil(c), (x, 100)); d.text((x, 12), t1, font=F(28), fill=(30, 30, 30)); d.text((x, 56), t2 + f"　{s}", font=F(22), fill=col)
    out.save(f"{OUT}/LOG463_04_old_vs_new_pick.jpg", quality=88)


# 06: 下ろし始め〜閉じた直後の上カメラ. 普段の本と爪が乗った本
def img06():
    rows = [("101157", "普段の本　下ろし始めから閉じるまで 2.1秒"), ("101434", "爪がブロックに乗った本　8.9秒"),
            ("102109", "位置2 の普段の本　2.1秒"), ("103504", "位置2 で爪が乗った本　8.9秒")]
    crops = {"101157": (330, 780, 380, 860), "101434": (330, 780, 380, 860), "102109": None, "103504": None}
    strips = []
    for s, title in rows:
        tp, tc = times(s); ts = np.linspace(tp - 0.6, tc + 0.3, 8); fr = []
        for t in ts:
            im = frame(s, t)
            if crops[s]: y0, y1, x0, x1 = crops[s]
            else:
                y0, y1, x0, x1 = 380, 830, 400, 880
            fr.append(cv2.resize(im[y0:y1, x0:x1], (240, 225)))
        strips.append((title, np.hstack(fr)))
    W = strips[0][1].shape[1] + 40; out = Image.new("RGB", (W, len(strips) * 290 + 20), "white"); d = ImageDraw.Draw(out)
    for i, (title, st) in enumerate(strips):
        y = 10 + i * 290; d.text((20, y + 4), title, font=F(26), fill=(30, 30, 30) if "普段" in title else (200, 30, 30))
        out.paste(pil(st), (20, y + 50))
    out.save(f"{OUT}/LOG463_06_snag_top.jpg", quality=86)


# 横からのスマホ動画のコマを並べる
def side(src, dst, title, n=12, cols=4):
    v = f"{SIDE}/{src}"
    dur = float(subprocess.check_output(["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "csv=p=0", v]).decode())
    subprocess.run(["ffmpeg", "-loglevel", "error", "-y", "-i", v, "-vf", f"fps={n}/{dur},scale=420:-1,tile={cols}x{n // cols}",
                    "-frames:v", "1", "/tmp/_side.jpg"], check=True)
    im = Image.open("/tmp/_side.jpg"); out = Image.new("RGB", (im.width, im.height + 60), "white")
    ImageDraw.Draw(out).text((10, 12), title, font=F(28), fill=(30, 30, 30)); out.paste(im, (0, 60))
    out.save(f"{OUT}/{dst}", quality=84)


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    img04(); img06()
    side("35e64696-IMG_4242-compressed.mp4", "LOG463_03_side_shallow.jpg", "前の教えた記録の再生を横から（左上から右へ時間順）。爪の先でブロックを挟んでいる")
    side("b2d7b71d-IMG_4246-compressed.mp4", "LOG463_07_side_snag.jpg", "位置2 で爪がブロックに乗った回を横から（左上から右へ時間順）。2段目で爪が乗って止まる")