#!/usr/bin/env python3
"""test_gap_fill.py — 実測角の抜けの埋め方 (build_correction_episodes --gap-fill. 依頼 46) の確かめ.

  python3 pi05/test_gap_fill.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


def samples():
    """0〜1 秒は 0.1 秒おきに止まって読めている → 1.0〜2.6 秒は送信で抜け (この間に j1 が 6° 動く)
    → 2.6〜3.6 秒はまた 0.1 秒おき."""
    a = np.round(np.arange(0.0, 1.01, 0.1), 3)
    b = np.round(np.arange(2.6, 3.61, 0.1), 3)
    at = np.concatenate([a, b])
    aq = np.zeros((len(at), 6))
    aq[len(a):, 1] = 6.0
    return at, aq


class TestFill(unittest.TestCase):
    def test_linear_is_interp(self):
        at, aq = samples()
        t = np.arange(0.0, 3.6, 0.1)
        q, gaps = bce.fill_motion_gaps(t, at, aq, mode="linear")
        self.assertTrue(np.allclose(q, np.stack([np.interp(t, at, aq[:, k]) for k in range(6)], axis=1)))
        self.assertEqual(len(gaps), 1)

    def test_speed20_moves_then_holds(self):
        at, aq = samples()
        t = np.arange(0.0, 3.6, 0.1)
        q, gaps = bce.fill_motion_gaps(t, at, aq, mode="speed20")
        self.assertAlmostEqual(gaps[0]["move_sec"], 6.0 / bce.SPEED20_DEG_PER_SEC)
        steps = np.abs(np.diff(q[:, 1]))
        self.assertLessEqual(steps.max(), bce.SPEED20_DEG_PER_SEC * 0.1 + 1e-6)   # 1 コマ 3° まで
        i = np.flatnonzero(t >= 1.25)[0]
        self.assertAlmostEqual(q[i, 1], 6.0)                                      # 0.2 秒で着いて止まる
        self.assertTrue(np.allclose(q[t <= 1.0], 0.0))                            # 抜けの前は実測のまま
        self.assertTrue(np.allclose(q[t >= 2.6, 1], 6.0))                            # 抜けのあとも実測のまま
        self.assertTrue(np.allclose(q[:, [0, 2, 3, 4, 5]], 0.0))                  # 動かない関節は動かさない

    def test_speed20_long_move_uses_whole_gap(self):
        at, aq = samples()
        aq[aq[:, 1] > 0, 1] = 90.0                     # 1.6 秒では速度 20 で着かない動き
        t = np.arange(0.0, 3.6, 0.1)
        q1, _ = bce.fill_motion_gaps(t, at, aq, mode="speed20")
        q0, _ = bce.fill_motion_gaps(t, at, aq, mode="linear")
        self.assertTrue(np.allclose(q1, q0))

    def test_unknown_mode(self):
        at, aq = samples()
        with self.assertRaises(ValueError):
            bce.fill_motion_gaps(np.arange(0, 1, 0.1), at, aq, mode="spline")


class TestMatch(unittest.TestCase):
    def test_match_correction_and_chunk(self):
        q0, q1, q2, q3 = (np.zeros(6), np.r_[0, 6, 0, 0, 0, 0.], np.r_[0, 12, 0, 0, 0, 0.],
                          np.r_[2, 12, 0, 0, 0, 0.])
        gaps = [{"ta": 1.0, "tb": 2.6, "move_deg": 6.0, "qb": q1},
                {"ta": 2.7, "tb": 4.3, "move_deg": 6.0, "qb": q2},
                {"ta": 10.0, "tb": 11.6, "move_deg": 2.0, "qb": q3},
                {"ta": 20.0, "tb": 23.6, "move_deg": 0.0, "qb": q3}]     # 爪 (動いていない)
        log = [{"chunk": 0, "t_exec_start": 0.9, "t_exec_end": 4.5,
                "sent_targets": [[0, 6.5, 0, 0, 0, 0], [0, 12, 0, 0, 0, 0]]},
               {"type": "correction", "action": "move", "sent": True, "t": 10.02, "after": [2, 12, 0, 0, 0, 0]},
               {"type": "correction", "action": "accept"}]
        m = bce.match_gaps_to_sends(gaps, log)
        self.assertEqual(m["motion_gaps"], 3)
        self.assertEqual(m["matched_chunk"], 2)
        self.assertEqual(m["matched_correction"], 1)
        self.assertEqual(m["unmatched"], 0)
        self.assertEqual(m["max_arrival_err_deg"], 0.5)

    def test_ambiguous_correction_not_paired(self):
        """codex 推奨: 補正の時刻が 2 つの抜けの候補に入るときは組にしない (差の集計から外す)."""
        gaps = [{"ta": 1.0, "tb": 2.6, "move_deg": 6.0, "qb": np.zeros(6)},
                {"ta": 2.7, "tb": 4.3, "move_deg": 6.0, "qb": np.ones(6) * 9}]
        log = [{"type": "correction", "action": "move", "sent": True, "t": 2.5, "after": [0] * 6}]
        m = bce.match_gaps_to_sends(gaps, log)
        self.assertEqual(m["matched_correction"], 0)
        self.assertEqual(m["ambiguous"], 1)
        self.assertIsNone(m["max_arrival_err_deg"])

    def test_chunk_count_mismatch(self):
        gaps = [{"ta": 1.0, "tb": 2.6, "move_deg": 6.0, "qb": np.zeros(6)}]
        log = [{"chunk": 0, "t_exec_start": 0.9, "t_exec_end": 4.5, "sent_targets": [[0] * 6, [0] * 6]}]
        m = bce.match_gaps_to_sends(gaps, log)
        self.assertEqual(m["chunk_count_mismatch"], 1)
        self.assertEqual(m["unmatched"], 1)


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