#!/usr/bin/env python3
"""test_rewind_correct.py — 巻き戻し・補正・補正エピソード作りの確かめ (アームなし. 2026-10-29).

  python3 pi05/test_rewind_correct.py          # 全部
  python3 pi05/test_rewind_correct.py -v

中身:
  - 巻き戻しの中継点の並び (逆順・始めの姿勢で終わる・移動量の上限で等分)
  - 補正の刻みと範囲チェック (apply_limits / check_send_safe を通る)
  - 補正エピソードの区間の除き方 (rewind と取り消した区間が入らない)
  - 合成の記録から build_correction_episodes がお手本と同じ形を作り, build_demo_episodes もそのまま読める
  - run_policy_client を偽の daemon / カメラ / モデルで通す: b / c の流れ (--dry-keys と --execute),
    b / c を押さないときは 変更前 (git HEAD) と同じログになる (項目が増えるだけ)
KIOXIA にもアームにも触らない (一時フォルダだけ).
"""
import argparse
import builtins
import importlib.util
import io
import json
import subprocess
import sys
import tempfile
import time
import types
import unittest
from pathlib import Path
from unittest import mock

import cv2
import numpy as np

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

import run_policy_client as rpc  # noqa: E402
import build_correction_episodes as bce  # noqa: E402
import build_demo_episodes as bde  # noqa: E402
from policy_recorder import STATES_COLS  # noqa: E402
from rewind_correct import (  # noqa: E402
    CORRECTION_STEPS, excluded_windows, next_step_index, parse_correction_key, rewind_waypoints, split_leg,
    step_keep_mask,
)

LO = np.array([-90.0, -10.0, -10.0, -30.0, -90.0, -60.0])
HI = np.array([90.0, 88.0, 60.0, 30.0, 60.0, 0.0])


class TestRewindWaypoints(unittest.TestCase):
    def test_reverse_and_end_at_start(self):
        start = [0, 0, 0, 0, 0, -30]
        sent = [[1, 0, 0, 0, 0, -30], [2, 1, 0, 0, 0, -30], [3, 2, 1, 0, 0, -30]]
        w = rewind_waypoints(sent, start)
        self.assertEqual(len(w), 3)
        np.testing.assert_allclose(w[0], sent[1])      # 最後に送った目標 (今いる所) の 1 つ手前から
        np.testing.assert_allclose(w[1], sent[0])
        np.testing.assert_allclose(w[-1], start)       # 始めの姿勢で終わる

    def test_single_target(self):
        w = rewind_waypoints([[5, 5, 5, 5, 5, -25]], [0, 0, 0, 0, 0, -30])
        self.assertEqual(len(w), 1)
        np.testing.assert_allclose(w[0], [0, 0, 0, 0, 0, -30])

    def test_split_by_max_move(self):
        start = np.zeros(6)
        sent = [np.array([30.0, 0, 0, 0, 0, 0]), np.array([45.0, 0, 0, 0, 0, 0])]
        w = rewind_waypoints(sent, start, current=sent[-1], max_move=20.0)
        # 45 → 30 (15°: 1 回) → 0 (30°: 2 回に等分)
        np.testing.assert_allclose([p[0] for p in w], [30.0, 15.0, 0.0])
        prev = sent[-1]
        for p in w:
            self.assertLessEqual(np.max(np.abs(p - prev)), 20.0 + 1e-9)
            prev = p
        np.testing.assert_allclose(w[-1], start)

    def test_split_leg_counts(self):
        self.assertEqual(len(split_leg(np.zeros(6), np.full(6, 20.0), 20.0)), 1)
        self.assertEqual(len(split_leg(np.zeros(6), np.full(6, 20.1), 20.0)), 2)


class TestCorrection(unittest.TestCase):
    def test_keys(self):
        self.assertEqual(parse_correction_key("1"), ("select", 0))
        self.assertEqual(parse_correction_key("6"), ("select", 5))
        self.assertEqual(parse_correction_key("+"), ("move", 1))
        self.assertEqual(parse_correction_key("."), ("move", 1))
        self.assertEqual(parse_correction_key("-"), ("move", -1))
        self.assertEqual(parse_correction_key(","), ("move", -1))
        self.assertEqual(parse_correction_key("\n"), ("accept", None))
        self.assertEqual(parse_correction_key("b"), ("cancel", None))
        self.assertEqual(parse_correction_key("g"), ("grip", None))
        self.assertIsNone(parse_correction_key("x"))
        self.assertEqual(CORRECTION_STEPS, (0.5, 1.0, 2.0))
        self.assertEqual(next_step_index(1, +1), 2)
        self.assertEqual(next_step_index(2, +1), 2)
        self.assertEqual(next_step_index(0, -1), 0)

    def test_step_size(self):
        cur = np.array([10.0, 20, 30, 0, -40, -30])
        for step in CORRECTION_STEPS:
            t, notes, unsafe = rpc.plan_correction_step(cur, 0, +1, step, cur, LO, HI, 20.0)
            self.assertEqual(unsafe, [])
            self.assertEqual(notes, [])
            np.testing.assert_allclose(t - cur, [step, 0, 0, 0, 0, 0])
            t, _, _ = rpc.plan_correction_step(cur, 3, -1, step, cur, LO, HI, 20.0)
            np.testing.assert_allclose(t - cur, [0, 0, 0, -step, 0, 0])

    def test_envelope_clamp(self):
        cur = np.array([89.5, 20, 30, 0, -40, -30])
        t, notes, unsafe = rpc.plan_correction_step(cur, 0, +1, 2.0, cur, LO, HI, 20.0)
        self.assertAlmostEqual(t[0], 90.0)               # 安全範囲で止まる (apply_limits)
        self.assertTrue(any("クランプ" in n for n in notes))
        self.assertEqual(unsafe, [])

    def test_unsafe_measured_outside(self):
        # 実測が範囲の外 (j1=110) だと, 移動量の制限で範囲の外へ戻る → check_send_safe が止める
        cur = np.array([0.0, 110, 30, 0, -40, -30])
        t, notes, unsafe = rpc.plan_correction_step(cur, 1, -1, 1.0, cur, LO, HI, 20.0)
        self.assertTrue(unsafe)

    def test_correction_session_dry(self):
        """補正モードを仮の腕で通す: j1 を +2 刻み, 刻み 2°, j0 -1, Enter."""
        arm = rpc.ArmIO(False, None, sim_q=[0, 0, 0, 0, 0, -30])
        log = io.StringIO()

        class Keys:
            def __init__(self, s): self.s = list(s)
            def get(self): return self.s.pop(0) if self.s else ""
            def flush(self): pass
        res, grip = rpc.run_correction(arm, Keys("2++]1-\n"), LO, HI, 20.0, 0, 1, 0, log, None)
        self.assertEqual(res["action"], "accept")
        np.testing.assert_allclose(res["final_target"], [-2.0, 2.0, 0, 0, 0, -30])
        lines = [json.loads(x) for x in log.getvalue().splitlines()]
        moves = [x for x in lines if x["action"] == "move"]
        self.assertEqual([m["joint"] for m in moves], [1, 1, 0])
        self.assertEqual([m["step_deg"] for m in moves], [1.0, 1.0, 2.0])
        self.assertTrue(all(m["type"] == "correction" for m in moves))
        self.assertEqual(moves[0]["before"], [0, 0, 0, 0, 0, -30])
        self.assertEqual(moves[0]["after"], [0, 1.0, 0, 0, 0, -30])

    def test_correction_cancel_returns(self):
        arm = rpc.ArmIO(False, None, sim_q=[0, 0, 0, 0, 0, -30])

        class Keys:
            def __init__(self, s): self.s = list(s)
            def get(self): return self.s.pop(0) if self.s else ""
            def flush(self): pass
        res, grip = rpc.run_correction(arm, Keys("1+++gb"), LO, HI, 20.0, 0, 1, 0, io.StringIO(), None)
        self.assertEqual(res["action"], "cancel")
        self.assertTrue(res["finished"])
        np.testing.assert_allclose(arm.read(), [0, 0, 0, 0, 0, -30])
        self.assertEqual(grip, 0)        # g で閉じたグリッパーも始めに戻す


class LaggyArm(rpc.ArmIO):
    """送っても半分しか動かない仮の腕 (未到達)."""
    def send(self, target, tag):
        t = np.asarray(target, dtype=float)
        self.sim_q = self.sim_q + (t - self.sim_q) * 0.5
        return False


class Keys:
    def __init__(self, s): self.s = list(s)
    def get(self): return self.s.pop(0) if self.s else ""
    def flush(self): pass


class TestAbnormal(unittest.TestCase):
    """codex R1 の指摘 1〜5 と A3 / C3 の異常系."""

    def test_rewind_not_reached_is_not_finished(self):
        arm = LaggyArm(False, None, sim_q=[10, 0, 0, 0, 0, -30])
        entry = {"kind": "chunk", "chunk": 0, "start_pose": np.array([0.0, 0, 0, 0, 0, -30]),
                 "sent_targets": [np.array([5.0, 0, 0, 0, 0, -30]), np.array([10.0, 0, 0, 0, 0, -30])],
                 "grip_changed": False, "t_start": 0.0, "t_end": 1.0}
        rw = rpc.run_rewind(entry, arm, LO, HI, 20.0, None)
        self.assertTrue(rw["traverse_completed"])
        self.assertFalse(rw["reached_start"])
        self.assertFalse(rw["finished"])           # 送り切っても戻れていなければ finished ではない

    def test_gripper_round_trip_counts(self):
        arm = rpc.ArmIO(False, None, sim_q=[0, 0, 0, 0, 0, -30])
        res, grip = rpc.run_correction(arm, Keys("1+gg\n"), LO, HI, 20.0, 0, 1, 0, io.StringIO(), None)
        self.assertEqual(res["grip_before"], res["grip_after"])
        self.assertTrue(res["grip_used"])          # 往復でも「動かした」 → 巻き戻さない
        res, grip = rpc.run_correction(arm, Keys("gg\n"), LO, HI, 20.0, 0, 2, 0, io.StringIO(), None)
        self.assertEqual(res["steps_sent"], 0)
        self.assertTrue(res["grip_used"])

    def test_correction_unreached_counts_from_measured(self):
        arm = LaggyArm(False, None, sim_q=[0, 0, 0, 0, 0, -30])
        res, _ = rpc.run_correction(arm, Keys("1]++\n"), LO, HI, 20.0, 0, 1, 0, io.StringIO(), None)
        # 1 刻み目 2° → 1° しか動かない → 2 刻み目は実測 1° から 2° = 3° へ (目標だけ 4° へ進まない)
        np.testing.assert_allclose(res["sent_targets"][1][0], 3.0)

    def test_cancel_unreached_keeps_gripper(self):
        class StuckOnReturn(rpc.ArmIO):
            """補正の刻みは届くが, 取り消しで戻るときは動かない腕."""
            def send(self, target, tag):
                if tag.startswith("pi05_corr_cancel"):
                    return False
                return super().send(target, tag)
        arm = StuckOnReturn(False, None, sim_q=[0, 0, 0, 0, 0, -30])
        with mock.patch("sys.stdout", new=io.StringIO()):
            res, grip = rpc.run_correction(arm, Keys("1]+++gb"), LO, HI, 20.0, 0, 1, 0, io.StringIO(), None)
        self.assertEqual(res["action"], "cancel")
        self.assertFalse(res["finished"])
        self.assertFalse(res["reached_start"])
        self.assertEqual(grip, 1)                  # 戻れていないのでグリッパーは閉じたまま

    def test_recorder_stop_timeout(self):
        import policy_recorder as pr

        class StuckCap:
            def __init__(self): self.ev = __import__("threading").Event()
            def read(self):
                self.ev.wait(); return False, None
        with tempfile.TemporaryDirectory() as td, mock.patch.object(pr, "STOP_JOIN_SEC", 0.3), \
                mock.patch("sys.stdout", new=io.StringIO()):
            cap = StuckCap()
            r = pr.PolicyRecorder(Path(td), cap, lambda: [0, 0, 0, 0, 0, -30], rpc.frame_to_model)
            r.start()
            time.sleep(0.2)
            r.stop()
            self.assertFalse(r.stopped_cleanly)
            self.assertTrue(r.alive())
            meta = json.loads((Path(td) / "recording" / "top_meta.json").read_text())
            self.assertEqual(meta["status"], "thread_not_stopped")
            cap.ev.set()


class TestImwriteFailure(unittest.TestCase):
    def test_imwrite_false_not_used(self):
        """JPEG の保存が失敗したら, そのエピソードは使わない (codex R3 P2)."""
        with tempfile.TemporaryDirectory() as td:
            td = Path(td)
            run = td / "runs" / "pi05_20990101_000200"
            write_fake_recording(run, [(0, 6, "model", 0, 30)])
            (run / "log.jsonl").write_text("")
            out = td / "ds" / "raw_episodes" / run.name
            with mock.patch.object(bce.cv2, "imwrite", lambda *a, **k: False):
                r = bce.build_one(run, out, argparse.Namespace(force=True, no_trim_mid_idle=False))
            self.assertFalse(r["used"])
            self.assertFalse(out.exists())


class TestExcludeSegments(unittest.TestCase):
    def test_rewind_and_cancel_excluded(self):
        # 0〜10 秒: model / 3〜5: 推論 3 を実行 (あとで巻き戻す) / 6〜7: rewind / 7〜8: model
        # 8〜9: correct (取り消し) 9〜9.5: rewind / 9.5〜: model
        t = np.arange(0, 12, 0.1)
        mode_ts = [0.0, 6.0, 7.0, 8.0, 9.0, 9.5]
        modes = ["model", "rewind", "model", "correct", "rewind", "model"]
        log = [{"chunk": 3, "executed": True, "t_exec_start": 3.0, "t_exec_end": 5.0},
               {"type": "rewind", "rewound": True, "chunk": 3, "segment_t": [3.0, 5.0],
                "t_start": 6.0, "t_end": 7.0, "finished": True},
               {"type": "correction", "action": "cancel", "t_start": 8.0, "t_end": 9.5}]
        w = excluded_windows(log)
        self.assertEqual(w, [(3.0, 7.0), (8.0, 9.5)])
        keep, m = step_keep_mask(t, mode_ts, modes, w)
        kt = t[keep]
        self.assertFalse(np.any((kt >= 3.0 - 1e-9) & (kt <= 7.0 + 1e-9)))
        self.assertFalse(np.any((kt >= 8.0 - 1e-9) & (kt <= 9.5 + 1e-9)))
        self.assertTrue(np.all(np.isin(m[keep], ["model", "correct"])))
        self.assertTrue(np.any(kt < 3.0) and np.any(kt > 9.5))

    def test_select_steps_joins(self):
        """巻き戻した区間の動きが入らず, 手前と続きがつながる (actions = つないだ並びの次の state)."""
        t = np.arange(0, 10, 0.1)
        q = np.zeros((len(t), 6))
        q[:, 5] = -30
        # 0〜2: j0 0→10, 2〜4: j0 10→40 (巻き戻す区間), 4〜6: rewind 40→10, 6〜10: j0 10→20 (j1 も)
        q[:, 0] = np.interp(t, [0, 2, 4, 6, 10], [0, 10, 40, 10, 20])
        q[:, 1] = np.interp(t, [0, 6, 10], [0, 0, 8])
        g = np.zeros(len(t))
        log = [{"type": "rewind", "rewound": True, "segment_t": [2.0, 4.0], "t_start": 4.0, "t_end": 6.0,
                "finished": True}]
        sel = bce.select_steps(t, q, g, np.array([0.0, 4.0, 6.0]), ["model", "rewind", "model"],
                               excluded_windows(log), trim_mid_idle=True)
        st = sel["state"]
        self.assertLessEqual(st[:, 0].max(), 20.0 + 1e-6)      # 40° まで行った区間が入っていない
        self.assertTrue(np.all(np.diff(t[sel["idx"]]) > 0))
        np.testing.assert_allclose(sel["actions"][:-1, :6], st[1:, :6])
        np.testing.assert_allclose(sel["actions"][-1, :6], st[-1, :6])
        # つなぎ目の跳びは通常の 1 歩の範囲 (10° 付近で手前と続きがつながる)
        self.assertLess(np.max(np.abs(np.diff(st[:, :6], axis=0))), 1.5)


def write_fake_recording(run_dir, segments, fps_cam=24.0, size=(640, 360), outcome="y", white=()):
    """合成の記録. segments = [(t0, t1, mode, j0_from, j0_to)]. 時刻は 1e9 秒台 (実時刻らしく)."""
    rec = run_dir / "recording"
    rec.mkdir(parents=True)
    base = 1.8e9
    t_end = segments[-1][1]
    with (rec / "states.csv").open("w") as f:
        f.write(",".join(STATES_COLS) + "\n")
        rows = []
        for (a, b, mode, j0a, j0b) in segments:
            rows.append((base + a, "stage", mode, None, None))
            for tt in np.arange(a, b, 0.1):
                j0 = j0a + (j0b - j0a) * (tt - a) / (b - a)
                rows.append((base + tt + 0.01, "actual", mode, [j0, 10, 5, 0, -40, -30], None))
        rows.append((base + 1.0, "gripper_cmd", "model", None, 0))
        rows.sort(key=lambda r: r[0])
        for ts, ev, mode, ang, gs in rows:
            angs = ["" if ang is None else f"{v:.3f}" for v in (ang or [None] * 6)]
            f.write(",".join([f"{ts:.6f}", ev, mode] + angs + ["", "" if gs is None else str(gs)] + [""] * 5) + "\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, t_end, 1.0 / fps_cam)):
            # white の窓 (除く区間) の中の画像だけ白 (250), ほかは暗い (k で少しずつ変える)
            v = 250 if any(a <= tt <= b for a, b in white) else 20 + (k % 40)
            img = np.full((size[1], size[0], 3), v, np.uint8)
            w.write(img)
            f.write(f"{k},{base + tt:.6f},model\n")
    w.release()
    (rec / "top_meta.json").write_text(json.dumps({"status": "ok"}))
    verdict = {"y": "yes", "n": "no"}[outcome]
    (run_dir / f"{run_dir.name}.json").write_text(json.dumps({
        "recording_enabled": True, "result": "completed",
        "outcome": {"y": "success", "n": "failure"}[outcome],
        "judgments": {"pick": {"human_verdict": verdict}, "place": {"human_verdict": verdict}}}))
    return base


class TestBuildCorrectionEpisodes(unittest.TestCase):
    def test_end_to_end(self):
        with tempfile.TemporaryDirectory() as td:
            td = Path(td)
            runs = td / "runs"
            good = runs / "pi05_20990101_000000"
            base = write_fake_recording(good, [
                (0, 2, "model", 0, 10), (2, 4, "model", 10, 10), (4, 6, "model", 10, 60),     # 4〜6 を巻き戻す
                (6, 8, "rewind", 60, 10), (8, 10, "correct", 10, 20), (10, 14, "model", 20, 35)],
                white=[(4.0, 8.0)])
            (good / "log.jsonl").write_text(json.dumps(
                {"type": "rewind", "rewound": True, "segment_t": [base + 4, base + 6],
                 "t_start": base + 6, "t_end": base + 8, "finished": True}) + "\n")
            bad = runs / "pi05_20990101_000100"
            write_fake_recording(bad, [(0, 6, "model", 0, 10)], outcome="n")
            (bad / "log.jsonl").write_text("")
            args = argparse.Namespace(force=True, no_trim_mid_idle=False)
            out = td / "ds" / "raw_episodes"
            r_bad = bce.build_one(bad, out / bad.name, args)
            self.assertFalse(r_bad["used"])                       # 失敗の run は使わない
            r = bce.build_one(good, out / good.name, args)
            self.assertTrue(r["used"], r)
            z = np.load(out / good.name / "episode.npz")
            meta = json.loads((out / good.name / "meta.json").read_text())
            self.assertLessEqual(z["state"][:, 0].max(), 35.0 + 1e-3)   # 60° まで行った区間がない
            self.assertLess(np.max(np.abs(np.diff(z["state"][:, :6], axis=0))), 3.0)   # つなぎ目が続く
            self.assertNotIn("rewind", meta["step_modes"])
            self.assertGreater(meta["correct_steps"], 0)
            self.assertEqual(meta["image_size"], [384, 216])
            frames = sorted((out / good.name / "frames").glob("*.jpg"))
            self.assertEqual(len(frames), len(z["state"]))
            # 除いた区間 (4〜8 秒) の画像は, 境界の最近傍としても入らない (codex R1-4)
            self.assertLess(max(float(cv2.imread(str(f)).mean()) for f in frames), 150)
            # お手本の作り方 (build_demo_episodes.build_one) もこの記録をそのまま読める
            ref = td / "ref" / "raw_episodes"
            rd = bde.build_one(good, ref / good.name, argparse.Namespace(
                include_failed=False, force=True, trim_mid_idle=True))
            self.assertTrue(rd["used"], rd)
            # 一緒に学習に使えるか (キー・形・FPS・画像の大きさ)
            self.assertEqual(bce.check_compat(out, ref), [])


# ---- run_policy_client を偽の daemon / カメラ / モデルで通す ---------------------------
class FakeWorld:
    def __init__(self):
        self.q = np.array([0.0, 0, 0, 0, 0, -30])
        self.sends = []
        self.grips = []
        self.events = []                # 送信とグリッパーを起きた順に (依頼 27: 下ろす → 開く → 戻る の順)
        self.fail_main_reads = 0        # 本体 (主スレッド) の読み取りをこの回数だけ失敗させる
        self.fail_after_label = None    # この送信のあとから失敗させる
        self.stuck_prefix = None        # この名前で始まる送信では動かない

    def get_angles(self, *a, **k):
        import threading
        if self.fail_main_reads and threading.current_thread() is threading.main_thread():
            self.fail_main_reads -= 1
            return None
        return [float(v) for v in self.q]

    def send_wait(self, target, label, tol, *a, **k):
        self.sends.append((label, [round(float(v), 3) for v in target]))
        self.events.append(("send", label, [round(float(v), 3) for v in target]))
        if self.stuck_prefix and label.startswith(self.stuck_prefix):
            return False, False         # この送信では動かない (未到達)
        self.q = np.array(target, dtype=float)
        if label == self.fail_after_label:
            self.fail_main_reads = 5    # read_angles の 5 回とも失敗 → 次の中継点を送らずに止まる
        return True, False


class FakeCap:
    def __init__(self, *a, **k): pass
    def isOpened(self): return True
    def grab(self): return True
    def read(self):
        time.sleep(0.01)
        return True, np.full((360, 640, 3), 100, np.uint8)
    def release(self): pass


class FakePolicy:
    """state から決まる行動 (画像は見ない): 毎ステップ j0 +1°, j1 +0.5°."""
    def __init__(self, *a, **k):
        self.uri, self.metadata = "fake", {}
    def infer(self, obs):
        s = np.asarray(obs["observation/state"], dtype=float)
        acts = np.zeros((10, 7))
        for k in range(10):
            acts[k, :6] = s[:6] + np.array([1.0, 0.5, 0, 0, 0, 0]) * (k + 1)
        return {"actions": acts}
    def close(self): pass


def run_client(mod, argv, inputs, td, keys="", fail_after_label=None, stuck_prefix=None, policy_cls=None):
    world = FakeWorld()
    world.fail_after_label = fail_after_label
    world.stuck_prefix = stuck_prefix
    raw = td / "raw"
    if not raw.exists():
        (raw / "ep0").mkdir(parents=True)
        st = np.zeros((20, 7), np.float32)
        st[:, 0] = np.linspace(-40, 40, 20); st[:, 1] = np.linspace(0, 60, 20)
        st[:, 4] = np.linspace(-80, 0, 20); st[:, 5] = -30
        np.savez(raw / "ep0" / "episode.npz", state=st, actions=st)
    log_root = td / "runs"
    it = iter(inputs)
    patches = [
        mock.patch.object(mod, "DATA_ROOT", td), mock.patch.object(mod, "LOG_ROOT", log_root),
        mock.patch.object(mod.camera_resolver, "resolve_cameras", lambda **k: {"top": 0}),
        mock.patch.object(mod.cv2, "VideoCapture", FakeCap),
        mock.patch.object(mod, "load_reference", lambda p: (np.zeros((216, 384, 3), np.uint8), "ref")),
        mock.patch.object(mod, "measure_framing", lambda a, b: (0.0, 0.0, 0.9)),
        mock.patch.object(mod, "PolicyClient", policy_cls or FakePolicy),
        mock.patch.object(mod, "get_daemon_socket", lambda *a, **k: None),
        mock.patch.object(mod, "close_daemon_socket", lambda *a, **k: None),
        mock.patch.object(mod, "daemon_get_angles", world.get_angles),
        mock.patch.object(mod, "daemon_set_gripper",
                          lambda s, *a, **k: (world.grips.append(s), world.events.append(("grip", s)))),
        mock.patch.object(mod, "daemon_set_gripper_value",
                          lambda v, *a, **k: (world.grips.append(("v", v)), world.events.append(("grip", ("v", v))))),
        mock.patch.object(mod, "_send_and_wait_arrival", world.send_wait),
        mock.patch.object(mod, "GRIPPER_SLEEP", 0.0),
        mock.patch.object(builtins, "input", lambda *a: next(it)),
        mock.patch.object(sys, "argv", ["run_policy_client.py", "--raw-dir", str(raw)] + argv),
        mock.patch.object(sys, "stdin", io.StringIO(keys)),
    ]
    for p in patches:
        p.start()
    try:
        with mock.patch("sys.stdout", new=io.StringIO()) as out:
            code = mod.main()
    finally:
        for p in reversed(patches):
            p.stop()
    run = sorted(log_root.glob("pi05_*"))[-1]
    lines = [json.loads(x) for x in (run / "log.jsonl").read_text().splitlines()]
    return code, world, run, lines, out.getvalue()


def strip(lines, drop=("t", "latency_s", "reference")):
    return [{k: v for k, v in d.items() if k not in drop} for d in lines]


class TestClientFlow(unittest.TestCase):
    def setUp(self):
        self.td = Path(tempfile.mkdtemp())
        time.sleep(1.1)       # run のフォルダ名は秒単位なので重ならないように

    def test_dry_keys_flow(self):
        code, world, run, lines, out = run_client(
            rpc, ["--dry-keys", "--steps-per-chunk", "10", "--via-stride", "3", "--max-chunks", "8"],
            ["", "", "b", "b", "c", "q"], self.td, keys="1+++\n")
        self.assertEqual(code, 0)
        self.assertEqual(world.sends, [])                     # アームには何も送っていない
        rw = [x for x in lines if x.get("type") == "rewind"]
        self.assertEqual([x["chunk"] for x in rw], [1, 0])    # 新しい方から 1 回分ずつさかのぼる
        self.assertTrue(all(x["rewound"] and x["finished"] for x in rw))
        np.testing.assert_allclose(rw[-1]["waypoints"][-1]["target"], [0, 0, 0, 0, 0, -30])
        ch = [x for x in lines if "chunk" in x and x.get("type") is None]
        self.assertEqual(len(ch), 6)                          # 推論の行はどれも消えていない
        corr = [x for x in lines if x.get("type") == "correction"]
        self.assertEqual([x["action"] for x in corr], ["move", "move", "move", "accept"])
        self.assertEqual(corr[-1]["final_target"][0], 3.0)
        # 補正のあとの推論 (最後の行) は補正した姿勢から
        self.assertAlmostEqual(ch[-1]["state"][0], 3.0, places=4)
        self.assertFalse((run / "recording").exists())       # ドライランは記録しない

    def test_execute_flow_with_recording(self):
        code, world, run, lines, out = run_client(
            rpc, ["--execute", "--steps-per-chunk", "10", "--via-stride", "3", "--settle-retries", "0",
                  "--state-from-target", "--ignore-framing", "--max-chunks", "8"],
            ["", "b", "c", "", "h", "y"], self.td, keys="1++\n")
        self.assertEqual(code, 0)
        labels = [s[0] for s in world.sends]
        self.assertEqual(labels[:4], ["pi05_chunk000_via02", "pi05_chunk000_via05",
                                      "pi05_chunk000_via08", "pi05_chunk000_via09"])
        rw_sends = [s for s in world.sends if s[0].startswith("pi05_rewind")]
        self.assertEqual([s[1][0] for s in rw_sends], [9.0, 6.0, 3.0, 0.0])   # 逆順で始めの姿勢へ
        self.assertEqual(rw_sends[-1][1], [0.0, 0.0, 0.0, 0.0, 0.0, -30.0])
        corr_sends = [s for s in world.sends if s[0].startswith("pi05_corr")]
        self.assertEqual([s[1][0] for s in corr_sends], [1.0, 2.0])
        self.assertEqual(labels[-1], "pi05_home")
        trial = json.loads((run / f"{run.name}.json").read_text())
        self.assertEqual(trial["outcome"], "success")
        self.assertEqual(trial["judgments"]["pick"]["human_verdict"], "yes")
        self.assertEqual(trial["end_reason"], "h")
        rec = run / "recording"
        for n in ("states.csv", "top.mp4", "top_frame_log.csv", "top_meta.json"):
            self.assertTrue((rec / n).exists(), n)
        stages = [l.split(",")[2] for l in (rec / "states.csv").read_text().splitlines()[1:]
                  if l.split(",")[1] == "stage"]
        # 入力を待っている間は wait (依頼 27). Enter / b / c のあとは model に戻ってから各状態へ
        self.assertEqual(stages, ["model", "wait", "model", "wait", "model", "rewind", "model",
                                  "wait", "model", "correct", "model", "wait", "model", "wait", "home"])
        ex = [x for x in lines if x.get("executed")]
        self.assertTrue(all("t_exec_start" in x and x["sent_targets"] for x in ex))

    def test_partial_chunk_can_be_rewound(self):
        """中継点の途中で止まった推論も history に残り, 次の b でそれが巻き戻される (codex R1-3)."""
        code, world, run, lines, out = run_client(
            rpc, ["--execute", "--steps-per-chunk", "10", "--via-stride", "3", "--settle-retries", "0",
                  "--ignore-framing", "--max-chunks", "6"],
            ["", "", "b", "q", ""], self.td, fail_after_label="pi05_chunk001_via02")
        rw = [x for x in lines if x.get("type") == "rewind"]
        self.assertEqual([x["chunk"] for x in rw], [1])          # 途中で止まった推論 1 を巻き戻す
        self.assertTrue(rw[0]["finished"])
        back = [s for s in world.sends if s[0].startswith("pi05_rewind")]
        self.assertEqual(back[-1][1], [10.0, 5.0, 0.0, 0.0, 0.0, -30.0])   # 推論 1 を始めた姿勢へ

    EXEC = ["--execute", "--steps-per-chunk", "10", "--via-stride", "3", "--settle-retries", "0",
            "--ignore-framing", "--max-chunks", "6"]

    def test_grip_in_correction_before_first_chunk(self):
        """初回推論の前に c → g で閉じたら, 開始時の開放をしない (codex R2-1)."""
        code, world, run, lines, out = run_client(rpc, self.EXEC, ["c", "", "q", ""], self.td, keys="g\n")
        self.assertEqual(world.grips[0], ("v", rpc.GRIP_VALUE))
        self.assertFalse(any(x.get("initial_gripper_open") for x in lines))
        ch = [x for x in lines if x.get("type") is None and "chunk" in x]
        self.assertEqual(ch[1]["state"][6], 1.0)              # モデルには「閉」を渡す

    def test_b_after_failed_rewind_is_refused(self):
        """巻き戻しが戻りきれなかったら, 次の b で古い区間へ飛ばない (codex R2-2)."""
        code, world, run, lines, out = run_client(rpc, self.EXEC, ["", "", "b", "b", "q", ""], self.td,
                                                  stuck_prefix="pi05_rewind")
        rw = [x for x in lines if x.get("type") == "rewind"]
        self.assertEqual([x["chunk"] for x in rw], [1])
        self.assertFalse(rw[0]["finished"])
        self.assertIn("これより前へは巻き戻しません", out)

    def test_b_after_failed_cancel_is_refused(self):
        code, world, run, lines, out = run_client(rpc, self.EXEC, ["", "c", "b", "q", ""], self.td,
                                                  keys="1]+++b", stuck_prefix="pi05_corr_cancel")
        self.assertEqual([x for x in lines if x.get("type") == "rewind"], [])
        cancel = [x for x in lines if x.get("action") == "cancel"]
        self.assertFalse(cancel[0]["finished"])
        self.assertFalse(any(s[0].startswith("pi05_rewind") for s in world.sends))
        self.assertIn("これより前へは巻き戻しません", out)

    def test_unchanged_without_b_c(self):
        """b / c を押さなければ 変更前 (git HEAD) と同じ送信・同じログ (項目が増えるだけ)."""
        src = subprocess.run(["git", "show", "HEAD:pi05/run_policy_client.py"], cwd=_HERE,
                             capture_output=True, text=True)
        if src.returncode != 0:
            self.skipTest("git HEAD の版が読めない")
        p = self.td / "run_policy_client_head.py"
        p.write_text(src.stdout)
        spec = importlib.util.spec_from_file_location("run_policy_client_head", p)
        old = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(old)
        for argv, inputs in (
                (["--execute", "--steps-per-chunk", "10", "--via-stride", "3", "--settle-retries", "0",
                  "--state-from-target", "--ignore-framing", "--max-chunks", "6"], ["", "", "s", "", "h", "y"]),
                (["--execute", "--max-chunks", "4"], ["", "", "", "q", ""]),
                (["--steps-per-chunk", "10", "--via-stride", "3", "--max-chunks", "3"], [])):
            c_old, w_old, _, l_old, _ = run_client(old, argv, inputs, self.td)
            time.sleep(1.1)
            c_new, w_new, _, l_new, _ = run_client(rpc, argv, inputs, self.td)
            time.sleep(1.1)
            self.assertEqual(c_old, c_new)
            self.assertEqual(w_old.sends, w_new.sends, argv)
            self.assertEqual(w_old.grips, w_new.grips, argv)
            self.assertEqual(len(l_old), len(l_new))
            new_keys = ("t_exec_start", "t_exec_end", "sent_targets")
            # HEAD が依頼 25 より前なら new_keys は新しい方にだけある. 依頼 25 以降なら両方にある (両方から外す)
            drop = ("t", "latency_s", "reference") + new_keys
            so, sn = strip(l_old, drop=drop), strip(l_new, drop=drop)
            self.assertEqual(so, sn, argv)


# ---- 依頼 27: つかんだ区間の巻き戻し (y / n) と wait のコマ ------------------------------
class GraspPolicy(FakePolicy):
    """状態で決まる仮のモデル. 開いていて j0 が 10〜30 → 3 ステップ目から閉じる (つかむ).
    閉じている間は j0 +1° / j1 -0.5° (持ち上げながら運ぶ). j0 が 30 以上 → 3 ステップ目から開く (離す)."""
    def infer(self, obs):
        s = np.asarray(obs["observation/state"], dtype=float)
        holding = s[6] > 0.5
        acts = np.zeros((10, 7))
        for k in range(10):
            acts[k, :6] = s[:6] + np.array([1.0, -0.5 if holding else 0.5, 0, 0, 0, 0]) * (k + 1)
            if holding:
                acts[k, 6] = 0.0 if (s[0] >= 30 and k >= 3) else 1.0
            else:
                acts[k, 6] = 1.0 if (10 <= s[0] < 30 and k >= 3) else 0.0
        return {"actions": acts}


class EventArm(rpc.ArmIO):
    """仮の腕. 送信とグリッパーを起きた順に残す."""
    def __init__(self, *a, **k):
        super().__init__(*a, **k)
        self.events = []

    def send(self, target, tag):
        self.events.append(("send", tag, [round(float(v), 3) for v in target]))
        return super().send(target, tag)

    def gripper(self, state):
        self.events.append(("grip", int(state)))
        return super().gripper(state)


class TestGraspRewind(unittest.TestCase):
    """つかんだ区間の巻き戻し (計算と送る順番). 腕は持ち上げた位置にいる."""
    S = np.array([10.0, 5, 0, 0, 0, -30])          # 推論を始めた姿勢
    A = np.array([12.0, 6, 0, 0, 0, -30])          # 途中の中継点
    G = np.array([13.0, 6.5, 0, 0, 0, -30])        # つかんだ位置 (閉じる前に最後に送った目標)
    UP = [13.0, 0.0, 0, 0, 0, -30]                 # 持ち上げた所 (今いる所)

    def entry(self):
        return {"kind": "chunk", "chunk": 1, "start_pose": self.S, "sent_targets": [self.A, self.G],
                "grip_changed": True, "grip_cmd": 1, "t_start": 0.0, "t_end": 1.0}

    def test_placed_order(self):
        arm = EventArm(False, None, sim_q=self.UP)
        with mock.patch("sys.stdout", new=io.StringIO()):
            rw = rpc.run_rewind(self.entry(), arm, LO, HI, 20.0, None, grasp="placed")
        # つかんだ位置まで下ろす → 開く → 逆にたどって始めの姿勢へ
        self.assertEqual(arm.events, [("send", "pi05_rewind_down_00", list(self.G)), ("grip", 0),
                                      ("send", "pi05_rewind_00", list(self.A)),
                                      ("send", "pi05_rewind_01", list(self.S))])
        self.assertEqual(rw["grasp_undo"], "placed")
        self.assertTrue(rw["gripper_opened"] and rw["reached_grasp"] and rw["finished"])

    def test_dropped_order(self):
        arm = EventArm(False, None, sim_q=self.UP)
        with mock.patch("sys.stdout", new=io.StringIO()):
            rw = rpc.run_rewind(self.entry(), arm, LO, HI, 20.0, None, grasp="dropped")
        # 今の位置ですぐ開く → 逆にたどって始めの姿勢へ (つかんだ位置へは下ろさない)
        self.assertEqual(arm.events, [("grip", 0), ("send", "pi05_rewind_00", list(self.A)),
                                      ("send", "pi05_rewind_01", list(self.S))])
        self.assertEqual(rw["grasp_undo"], "dropped")
        self.assertTrue(rw["gripper_opened"] and rw["finished"])

    def test_placed_not_reached_keeps_closed(self):
        """つかんだ位置まで下ろせなければ開かない (宙で落とさない) → 戻りきれない扱い."""
        class Stuck(EventArm):
            def send(self, target, tag):
                self.events.append(("send", tag, [round(float(v), 3) for v in target]))
                return False                     # 動かない
        arm = Stuck(False, None, sim_q=self.UP)
        with mock.patch("sys.stdout", new=io.StringIO()):
            rw = rpc.run_rewind(self.entry(), arm, LO, HI, 20.0, None, grasp="placed")
        self.assertNotIn(("grip", 0), arm.events)
        self.assertFalse(rw["gripper_opened"])
        self.assertFalse(rw["finished"])


class TestGraspClientFlow(unittest.TestCase):
    """run_policy_client を偽の daemon で通す. 推論 0 = 近づく / 1 = つかむ / 2 = 運ぶ."""
    EXEC = ["--execute", "--steps-per-chunk", "10", "--via-stride", "3", "--settle-retries", "0",
            "--ignore-framing", "--max-chunks", "8"]

    def setUp(self):
        self.td = Path(tempfile.mkdtemp())
        time.sleep(1.1)

    def run_(self, inputs, argv=None):
        return run_client(rpc, argv or self.EXEC, inputs, self.td, policy_cls=GraspPolicy)

    @staticmethod
    def after_carry_rewind(world):
        """推論 2 (運ぶ) を巻き戻し終えたあとの出来事."""
        ev = world.events
        last = max(k for k, e in enumerate(ev) if e[0] == "send" and e[1].startswith("pi05_rewind")
                   and e[2][0] == 13.0)
        return ev[:last + 1], ev[last + 1:]

    def test_carry_rewind_keeps_gripper_closed(self):
        """運ぶ途中の区間の b: グリッパーは閉じたまま, 送った目標を逆の順にたどる (白箱にぶつかったとき)."""
        code, world, run, lines, out = self.run_(["", "", "", "b", "q", ""])
        self.assertEqual(code, 0)
        carry = [s[1][0] for s in world.sends if s[0].startswith("pi05_chunk002")]
        self.assertEqual(carry, [16.0, 19.0, 22.0, 23.0])
        back = [s[1][0] for s in world.sends if s[0].startswith("pi05_rewind")]
        self.assertEqual(back, [22.0, 19.0, 16.0, 13.0])         # 逆の順で推論 2 の始め (つかんだ位置) へ
        self.assertEqual(world.grips, [0, ("v", rpc.GRIP_VALUE)])  # 開始時の開放と つかむ だけ. 開いていない
        rw = [x for x in lines if x.get("type") == "rewind"]
        self.assertEqual([x["chunk"] for x in rw], [2])
        self.assertNotIn("grasp_undo", rw[0])
        self.assertTrue(rw[0]["finished"])
        ch = [x for x in lines if x.get("type") is None and "chunk" in x]
        self.assertEqual(ch[-1]["state"][6], 1.0)                 # 戻ったあともモデルには「閉」

    def test_grasp_rewind_y(self):
        """運ぶ区間を戻したあと, つかんだ区間の b で y: 下ろす → 開く → 始めの姿勢へ."""
        code, world, run, lines, out = self.run_(["", "", "", "b", "b", "y", "q", ""])
        self.assertEqual(code, 0)
        _, tail = self.after_carry_rewind(world)
        # つかんだ位置 (13, 6.5) には運ぶ区間の巻き戻しで既に着いているので, 下ろす送信は 0.5° 以下で省く
        self.assertEqual(tail, [("grip", 0), ("send", "pi05_rewind_00", [10.0, 5.0, 0, 0, 0, -30.0])])
        rw = [x for x in lines if x.get("type") == "rewind"]
        self.assertEqual([x["chunk"] for x in rw], [2, 1])
        self.assertEqual(rw[1]["grasp_undo"], "placed")
        self.assertTrue(rw[1]["reached_grasp"] and rw[1]["gripper_opened"] and rw[1]["finished"])
        self.assertEqual(rw[1]["grasp_pose"], [13.0, 6.5, 0, 0, 0, -30.0])
        self.assertEqual([w["leg"] for w in rw[1]["waypoints"]], ["down", "back"])
        ch = [x for x in lines if x.get("type") is None and "chunk" in x]
        self.assertEqual(ch[-1]["state"][:2], [10.0, 5.0])        # 推論 1 の始めの姿勢から
        self.assertEqual(ch[-1]["state"][6], 0.0)                 # 開いた状態で再推論
        self.assertNotIn("置き直したら Enter", out)

    def test_grasp_rewind_n(self):
        """つかんだ区間の b で n: 開く → 始めの姿勢へ → 写真のパスを出して Enter を待つ."""
        code, world, run, lines, out = self.run_(["", "", "", "b", "b", "n", "", "q", ""])
        self.assertEqual(code, 0)
        _, tail = self.after_carry_rewind(world)
        self.assertEqual(tail, [("grip", 0), ("send", "pi05_rewind_00", [10.0, 5.0, 0, 0, 0, -30.0])])
        rw = [x for x in lines if x.get("type") == "rewind"]
        self.assertEqual(rw[1]["grasp_undo"], "dropped")
        self.assertTrue(rw[1]["gripper_opened"] and rw[1]["finished"])
        self.assertEqual([w.get("leg") for w in rw[1]["waypoints"]], [None])
        frame = str(run / "frames" / "0001.jpg")
        self.assertEqual(rw[1]["start_frame"], frame)
        self.assertIn(frame, out)
        self.assertTrue(Path(frame).exists())
        # 置き直す間も wait
        stages = [l.split(",")[2] for l in (run / "recording" / "states.csv").read_text().splitlines()[1:]
                  if l.split(",")[1] == "stage"]
        k = len(stages) - 1 - stages[::-1].index("rewind")
        self.assertEqual(stages[k:k + 4], ["rewind", "model", "wait", "model"])

    def test_grasp_rewind_other_cancels(self):
        code, world, run, lines, out = self.run_(["", "", "", "b", "b", "x", "q", ""])
        rw = [x for x in lines if x.get("type") == "rewind"]
        self.assertEqual([x["chunk"] for x in rw], [2])           # つかんだ区間は戻していない
        self.assertEqual(world.grips, [0, ("v", rpc.GRIP_VALUE)])
        self.assertIn("やめました", out)

    def test_release_refused(self):
        """離した区間は今のまま断る. 置き直すなら h で終えてやり直す, と出す."""
        code, world, run, lines, out = self.run_(["", "", "", "", "", "b", "q", ""])
        self.assertEqual(world.grips, [0, ("v", rpc.GRIP_VALUE), 0])   # 推論 4 で離した
        self.assertEqual([x for x in lines if x.get("type") == "rewind"], [])
        self.assertIn("ブロックを置き直すなら h で終えてやり直してください", out)

    def test_dry_keys_grasp_y(self):
        code, world, run, lines, out = run_client(
            rpc, ["--dry-keys", "--steps-per-chunk", "10", "--via-stride", "3", "--max-chunks", "8"],
            ["", "", "", "b", "b", "y", "q"], self.td, policy_cls=GraspPolicy)
        self.assertEqual(world.sends, [])
        rw = [x for x in lines if x.get("type") == "rewind"]
        self.assertEqual([x["chunk"] for x in rw], [2, 1])
        self.assertEqual(rw[1]["grasp_undo"], "placed")
        self.assertTrue(rw[1]["finished"])


class TestWaitExcluded(unittest.TestCase):
    def test_wait_rewind_and_undone_not_in_episode(self):
        """wait のコマ・rewind のコマ・巻き戻しで取り消した区間は, 停止の詰め方 (先頭を残す) を
        通っても学習用エピソードに入らない."""
        with tempfile.TemporaryDirectory() as td:
            td = Path(td)
            run = td / "runs" / "pi05_20990101_000300"
            # 0〜3 近づく / 3〜5 つかんで運ぶ (あとで n で巻き戻す) / 5〜7 rewind / 7〜10 wait (置き直す)
            # / 10〜14 model. 3〜10 秒の画像は白 (手が写る想定)
            base = write_fake_recording(run, [
                (0, 3, "model", 0, 15), (3, 5, "model", 15, 40), (5, 7, "rewind", 40, 15),
                (7, 10, "wait", 15, 15), (10, 14, "model", 15, 30)], white=[(3.0, 9.99)])
            (run / "log.jsonl").write_text(json.dumps(
                {"type": "rewind", "rewound": True, "grasp_undo": "dropped", "gripper_opened": True,
                 "segment_t": [base + 3, base + 5], "t_start": base + 5, "t_end": base + 7,
                 "finished": True}) + "\n")
            out = td / "ds" / "raw_episodes" / run.name
            r = bce.build_one(run, out, argparse.Namespace(force=True, no_trim_mid_idle=False))
            self.assertTrue(r["used"], r)
            meta = json.loads((out / "meta.json").read_text())
            z = np.load(out / "episode.npz")
            self.assertEqual(set(meta["step_modes"]), {"model"})
            self.assertGreater(meta["excluded_mode_steps"], 0)       # rewind と wait
            self.assertGreater(meta["excluded_window_steps"], 0)     # 取り消した区間
            self.assertLessEqual(z["state"][:, 0].max(), 30.0 + 1e-3)
            self.assertLess(np.max(np.abs(np.diff(z["state"][:, :6], axis=0))), 3.0)
            frames = sorted((out / "frames").glob("*.jpg"))
            self.assertEqual(len(frames), len(z["state"]))
            self.assertLess(max(float(cv2.imread(str(f)).mean()) for f in frames), 150)

    def test_select_steps_drops_long_wait(self):
        """止まっている区間の詰め方は先頭 2 コマを残すが, wait は詰める前に外すので 1 コマも残らない."""
        t = np.arange(0, 10, 0.1)
        q = np.zeros((len(t), 6)); q[:, 5] = -30
        q[:, 0] = np.interp(t, [0, 3, 7, 10], [0, 10, 10, 20])     # 3〜7 秒は止まっている
        g = np.zeros(len(t))
        sel = bce.select_steps(t, q, g, np.array([0.0, 3.0, 7.0]), ["model", "wait", "model"], [],
                               trim_mid_idle=True)
        kt = t[sel["idx"]]
        self.assertFalse(np.any((kt >= 3.0 - 1e-9) & (kt < 7.0 - 1e-9)))
        self.assertNotIn("wait", sel["modes"])
        self.assertEqual(sel["counts"]["excluded_mode"], 40)


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