#!/usr/bin/env python3
"""test_trim_hold.py — 着いたあとの止まりの詰め方 (build_correction_episodes --trim-hold. 依頼 48) の確かめ.

  python3 pi05/test_trim_hold.py -v

ファイルを読まない部分だけ (KIOXIA にもアームにも触らない).
"""
import sys
import unittest
from pathlib import Path

import numpy as np

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

import build_correction_episodes as bce  # noqa: E402
import build_demo_episodes as bde  # noqa: E402


def episode(pattern):
    """pattern の 1 文字 = 1 コマ. '#' = j1 が 2° 動く / '.' = 止まる / 'g' = 爪を閉じる指令のコマ (止まったまま).
    戻り値: state, actions (= 次のコマの state. 最後は自分自身. 爪は次の指令)."""
    n = len(pattern) + 1
    s = np.zeros((n, 7))
    for i, c in enumerate(pattern):
        s[i + 1] = s[i]
        if c == "#":
            s[i + 1, 1] += 2.0
        if c == "g":
            s[i + 1, 6] = 1.0
    a = np.concatenate([s[1:], s[-1:]])
    return s, a


def idle_of(s, a):
    return (np.max(np.abs(a[:, :6] - s[:, :6]), axis=1) < bde.MID_IDLE_EPS_DEG) & (a[:, 6] == s[:, 6])


class TestHoldKeep(unittest.TestCase):
    def test_same_as_trim2_when_all_two(self):
        """残す数が全部 MID_IDLE_KEEP_STEPS なら trim2 の mid_idle_keep_mask と同じ."""
        rng = np.random.default_rng(0)
        k = bde.MID_IDLE_KEEP_STEPS
        for _ in range(300):
            n = 80
            s = np.cumsum(rng.choice([0, 0, 0, 0.1, 0.2, 1.5], size=(n, 7)), axis=0)
            s[:, 6] = np.arange(n) > rng.integers(5, 75)
            a = np.concatenate([s[1:], s[-1:]])
            modes = list(rng.choice([bce.MODE_MODEL, bce.MODE_CORRECT], size=n))
            np.testing.assert_array_equal(bce.hold_keep_mask(s, a, modes, k, k, k), bde.mid_idle_keep_mask(s, a))

    def test_model_holds_removed_correct_keeps_one(self):
        s, a = episode("##....##....##....#....#....##")
        n = len(s)
        modes = [bce.MODE_MODEL] * 14 + [bce.MODE_CORRECT] * (n - 14)
        keep = bce.hold_keep_mask(s, a, modes)
        idle = idle_of(s, a)
        regions, i = [], 0
        while i < n:
            if idle[i]:
                j = i
                while j < n and idle[j]:
                    j += 1
                regions.append((i, j))
                i = j
            else:
                i += 1
        for i, j in regions[:-1]:          # 最後の区間はエピソードの終わり (最後のコマは必ず残す)
            want = 1 if modes[i] == bce.MODE_CORRECT else 0
            self.assertEqual(int(keep[i:j].sum()), want, (i, j, modes[i]))
        self.assertTrue(keep[~idle].all())                     # 動いているコマは消さない
        self.assertTrue(keep[0] and keep[-1])

    def test_gripper_neighbors_keep_like_trim2(self):
        """爪の切り替えのすぐ前・すぐ後の止まりは trim2 と同じ先頭 2 コマ (+ 切り替えの前後 1 コマ)."""
        s, a = episode("##......g......##......##")
        modes = [bce.MODE_MODEL] * len(s)
        keep = bce.hold_keep_mask(s, a, modes)
        trim2 = bde.mid_idle_keep_mask(s, a)
        g = int(np.flatnonzero(a[:, 6] != s[:, 6])[0])     # 指令のコマ. 切り替わるのは g + 1
        np.testing.assert_array_equal(keep[:g + 9], trim2[:g + 9])     # 爪の前後は trim2 と同じ
        self.assertEqual(int(keep[2:g].sum()), bde.MID_IDLE_KEEP_STEPS)   # 直前の区間: 先頭 2
        self.assertTrue(keep[g:g + 3].all())                              # 切り替え (g + 1) の前後 1 コマ
        # 直後の区間 (g + 1 〜 g + 6) は先頭 2 コマ (g + 1, g + 2. 前後の保護と重なる) だけ
        self.assertEqual(int(keep[g + 1:g + 7].sum()), bde.MID_IDLE_KEEP_STEPS)
        # 爪から離れた止まり (model) は残さない
        far = idle_of(s, a)
        far[:g + 9] = False
        far[-1] = False
        self.assertFalse(keep[far].any())

    def test_slow_drift_not_trimmed(self):
        """0.2° ずつ動き続ける区間はずれで区切られ, model (残す数 0) でも 0.5° ごとの姿勢が残る (codex 必須)."""
        n = 30
        s = np.zeros((n, 7))
        s[:, 1] = np.arange(n) * 0.2
        a = np.concatenate([s[1:], s[-1:]])
        for mode in (bce.MODE_MODEL, bce.MODE_CORRECT):
            keep = bce.hold_keep_mask(s, a, [mode] * n)
            jump = np.abs(np.diff(s[keep, 1])).max()
            self.assertLessEqual(jump, bde.MID_IDLE_DRIFT_DEG + 0.2 + 1e-9, mode)   # 詰めても 1 コマ 0.7° まで
            self.assertGreaterEqual(int(keep.sum()), n // 3, mode)
        self.assertGreaterEqual(int(bce.hold_keep_mask(s, a, [bce.MODE_CORRECT] * n).sum()),
                                int(bce.hold_keep_mask(s, a, [bce.MODE_MODEL] * n).sum()))

    def test_select_steps_trim_hold(self):
        """select_steps(trim_hold=True) で止まりが減り, actions が「次のコマの state」のまま."""
        t = np.arange(0, 6.0, 0.1)
        q = np.zeros((len(t), 6))
        for k in range(len(t)):            # 0.5 秒ごとに 1 コマで 2° 動いて止まる
            q[k, 1] = 2.0 * (k // 5)
        g = np.where(t >= 3.0, 1.0, 0.0)
        mode_ts, modes = np.array([0.0]), [bce.MODE_MODEL]
        old = bce.select_steps(t, q, g, mode_ts, modes, [], trim_mid_idle=True)
        new = bce.select_steps(t, q, g, mode_ts, modes, [], trim_mid_idle=True, trim_hold=True)
        self.assertLess(len(new["idx"]), len(old["idx"]))
        self.assertLess(idle_of(new["state"], new["actions"]).mean(), idle_of(old["state"], old["actions"]).mean())
        np.testing.assert_array_equal(new["actions"][:-1, :6], new["state"][1:, :6])
        self.assertEqual(new["state"][-1, 1], q[-1, 1])                   # 終わりの姿勢は変わらない
        self.assertEqual(int(np.sum(np.diff(new["state"][:, 6]) != 0)), 1)   # 爪の切り替えは残る


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