#!/usr/bin/env python3
"""test_pad_build.py — ゲームパッドの補正 (run_policy_client の p. cd2fb46) の記録から
build_correction_episodes.py が学習データを作れるかの確かめ (2026-11-06 依頼 97).

  python3 pi05/test_pad_build.py -v

合成の記録だけ (KIOXIA にもアームにも触らない). 確かめること:
  - 補正の間 (状態 correct) のコマが学習に入る. 長押しで動いている途中のコマも入る
  - 補正の中で送った爪の指令 (recorder.gripper → states.csv の gripper_cmd 行) の時刻で 0/1 が切り替わる
  - 補正を始めたときの約 3 秒の抜け (ゲームパッドの窓を開く間. 腕は止まっている) と, 長押しの途中の 0.3 秒の抜けは,
    speed20 で埋めても linear と同じ (抜けが短く, 速度 20 の速さでは抜けいっぱいを使う)
  - 長押しを離して stop で止まったところは, 今までと同じ決まり (--trim-hold. correct は 1 コマ) で詰まる
"""
import argparse
import json
import sys
import tempfile
import unittest
from pathlib import Path

import cv2
import numpy as np

_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE))

import build_correction_episodes as bce  # noqa: E402

COLS = ["ts", "event_type", "stage", "j0", "j1", "j2", "j3", "j4", "j5", "speed", "gripper_state",
        "ts_send", "ts_arrival", "attempt", "arrived", "note"]
BASE = 1.8e9


def pose_at(t):
    """合成の動き (秒 → 6 関節).
    0〜2 秒 model: j0 0 → 10 / 2〜5 秒 correct: 止まり (ゲームパッドの窓を開く間. 実測なし)
    5.0〜5.6 秒 correct: j1 を長押しで 10 → 46 (60°/秒. 5.2〜5.5 秒は実測なし) / 5.6 秒 stop
    5.6〜9 秒 correct: 止まり (8.0 秒に爪を閉じる指令) / 9〜11 秒 model: j0 10 → 20"""
    j0 = 10.0 * min(t, 2.0) / 2.0 + (10.0 * (t - 9.0) / 2.0 if t > 9.0 else 0.0)
    j1 = 10.0 + 60.0 * min(max(t - 5.0, 0.0), 0.6)
    return [j0, j1, 5.0, 0.0, -40.0, -30.0]


def sample_times():
    t = np.round(np.arange(0.0, 11.0, 0.1), 3)
    gap = ((t > 2.05) & (t < 5.0)) | ((t > 5.15) & (t < 5.5))
    return t[~gap]


def write_run(run_dir, size=(640, 360), fps_cam=10.0):
    rec = run_dir / "recording"
    rec.mkdir(parents=True)
    rows = [(0.0, "stage", "model", None, None, ""), (2.0, "stage", "correct", None, None, ""),
            (9.0, "stage", "model", None, None, ""),
            (0.5, "gripper_cmd", "model", None, 0, ""),
            (8.0, "gripper_cmd", "correct", None, 1, ""),                        # 補正の中の爪 (recorder.gripper)
            (8.0001, "rpc", "correct", None, None, "cmd=set_gripper_value rpc_sec=0.013"),
            (5.0, "rpc", "correct", None, None, "cmd=send_angle rpc_sec=0.034"),
            (5.6, "rpc", "correct", None, None, "cmd=stop rpc_sec=0.025")]
    for t in sample_times():
        rows.append((float(t) + 0.01, "actual", "model" if t < 2.0 or t >= 9.0 else "correct", pose_at(t), None, ""))
    rows.sort(key=lambda r: r[0])
    with (rec / "states.csv").open("w") as f:
        f.write(",".join(COLS) + "\n")
        for ts, ev, mode, ang, gs, note in rows:
            angs = ["" if ang is None else f"{v:.3f}" for v in (ang or [None] * 6)]
            f.write(",".join([f"{BASE + ts:.6f}", ev, mode] + angs + ["", "" if gs is None else str(gs)]
                             + [""] * 4 + [note]) + "\n")
    w = cv2.VideoWriter(str(rec / "top.mp4"), cv2.VideoWriter_fourcc(*"mp4v"), fps_cam, size)
    with (rec / "top_frame_log.csv").open("w") as f:
        f.write("frame_idx,capture_ts,mode\n")
        for k, tt in enumerate(np.arange(0, 11.0, 1.0 / fps_cam)):
            w.write(np.full((size[1], size[0], 3), 20 + (k % 40), np.uint8))
            f.write(f"{k},{BASE + tt:.6f},model\n")
    w.release()
    (rec / "top_meta.json").write_text(json.dumps({"status": "ok"}))
    (run_dir / f"{run_dir.name}.json").write_text(json.dumps({
        "recording_enabled": True, "result": "completed", "outcome": "success",
        "judgments": {"pick": {"human_verdict": "yes"}, "place": {"human_verdict": "yes"}}}))
    log = [{"type": "pad", "correction_id": 1, "kind": "move", "sent": True, "t_send": BASE + 5.0, "joint": 1,
            "no_check": True, "command": [10, 25, 5, 0, -40, -30]},
           {"type": "pad", "correction_id": 1, "kind": "halt", "t_halt": BASE + 5.6, "halt": "stop", "halt_joint": 1,
            "measured_after_halt": [10, 46, 5, 0, -40, -30]},
           {"type": "pad", "correction_id": 1, "kind": "grip_close", "sent": True, "grip_cmd": 1},
           {"type": "correction", "action": "accept", "device": "gamepad", "correction_id": 1, "after_chunk": 0,
            "t_start": BASE + 2.0, "t_end": BASE + 9.0, "grip_used": True}]
    (run_dir / "log.jsonl").write_text("".join(json.dumps(x) + "\n" for x in log))


class TestPadBuild(unittest.TestCase):
    def build(self, gap_fill):
        td = Path(tempfile.mkdtemp())
        run = td / "runs" / "pi05_20991106_000000"
        write_run(run)
        out = td / "ds" / "raw_episodes" / run.name
        args = argparse.Namespace(force=True, no_trim_mid_idle=False, gap_fill=gap_fill, trim_hold=True)
        res = bce.build_one(run, out, args)
        self.assertTrue(res["used"], res)
        z = np.load(out / "episode.npz")
        meta = json.loads((out / "meta.json").read_text())
        return z, meta

    def test_correct_steps_and_gripper(self):
        z, meta = self.build("speed20")
        st, modes, t = z["state"], meta["step_modes"], z["t_orig"]
        self.assertGreater(meta["correct_steps"], 0)
        # 長押しで動いている途中のコマ (j1 が 10 と 46 の間) が correct として入る
        mid = [i for i in range(len(st)) if 12.0 < st[i, 1] < 44.0]
        self.assertGreaterEqual(len(mid), 3)
        self.assertTrue(all(modes[i] == "correct" for i in mid))
        # 爪は補正の中の指令の時刻で 0 → 1 に切り替わる (切り替わる前後のコマは correct)
        flip = int(np.flatnonzero(np.diff(st[:, 6]) != 0)[0])
        self.assertEqual(st[flip, 6], 0.0)
        self.assertEqual(st[flip + 1, 6], 1.0)
        self.assertEqual(modes[flip + 1], "correct")
        self.assertEqual(int(np.sum(np.diff(st[:, 6]) != 0)), 1)
        self.assertTrue(np.all(np.diff(t) > 0))

    def test_stop_after_release_is_trimmed(self):
        z, meta = self.build("speed20")
        st, a, modes = z["state"], z["actions"], meta["step_modes"]
        # 止まり = 関節が動かず爪も切り替わらないコマ (check_dataset と同じ. 爪を切り替えるコマは数えない)
        still = (np.max(np.abs(a[:, :6] - st[:, :6]), axis=1) < 0.3) & (a[:, 6] == st[:, 6])
        run, longest = 0, 0
        for i in range(len(st)):
            run = run + 1 if still[i] and modes[i] == "correct" else 0
            longest = max(longest, run)
        # 長押しを離したあとの止まり (5.6〜9 秒) は, 爪の切り替えの直前・直後の 2 コマずつ (trim2 と同じ) まで
        self.assertLessEqual(longest, 2)
        held = [i for i in range(len(st)) if abs(st[i, 1] - 46.0) < 1e-3 and modes[i] == "correct"]
        self.assertLessEqual(len(held), 2 + 2 + 1)        # correct の 1 + 爪の切り替えの前後 (keep_grip 2 + guard)

    def test_speed20_same_as_linear_for_pad_gaps(self):
        z1, m1 = self.build("speed20")
        z2, m2 = self.build("linear")
        np.testing.assert_allclose(z1["state"], z2["state"])
        self.assertEqual(m1["step_modes"], m2["step_modes"])
        self.assertEqual(m1["gap_fill"]["gaps"], 2)                 # 窓を開く間の 3 秒と, 長押しの途中の 0.4 秒
        self.assertEqual(m1["gap_fill"]["motion_gaps"], 1)


if __name__ == "__main__":
    unittest.main()
