#!/usr/bin/env python3
"""test_state_source.py — 1 コマごとの state の出どころ (state_source. 依頼 196) の確かめ.

  python3 pi05/test_state_source.py -v

前半はファイルを読まない部分だけ. 後半 (TestSameAsBefore) は KIOXIA の runs/pi05_policy を読み,
直す前のコード (git の HEAD) と今のコードで /tmp に作り, state / actions / t がバイト単位で同じことを見る
(KIOXIA がなければ飛ばす. データセットにもアームにも触らない).
"""
import importlib.util
import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace

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(still=False):
    """0〜1 秒は 0.1 秒おき → 1.0〜2.6 秒は抜け (j1 が 6° 動く. still なら動かない) → 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))
    if not still:
        aq[len(a):, 1] = 6.0
    return at, aq


def runs_of_nonzero(src):
    nz = np.r_[0, (np.asarray(src) != 0).astype(int)]
    return int((np.diff(nz) == 1).sum())


class TestFlags(unittest.TestCase):
    def test_motion_gap_is_1(self):
        at, aq = samples()
        t = np.arange(0.0, 3.6, 0.1)
        src = bce.gap_fill_flags(t, at, aq)
        self.assertEqual(src.dtype, np.int8)
        inside = (t > 1.0) & (t < 2.6)
        self.assertTrue(np.all(src[inside] == bce.SRC_FILLED_MOVE))
        self.assertTrue(np.all(src[~inside] == bce.SRC_MEASURED))

    def test_still_gap_is_2(self):
        at, aq = samples(still=True)
        t = np.arange(0.0, 3.6, 0.1)
        src = bce.gap_fill_flags(t, at, aq)
        self.assertEqual(set(np.unique(src)), {bce.SRC_MEASURED, bce.SRC_FILLED_STILL})

    def test_flags_mark_exactly_what_speed20_changes(self):
        """speed20 が線形補間から変えたコマは, すべて印が 0 でないコマ."""
        at, aq = samples()
        t = np.arange(0.0, 3.6, 0.1)
        q_lin, _ = bce.fill_motion_gaps(t, at, aq, mode="linear")
        q_sp, _ = bce.fill_motion_gaps(t, at, aq, mode="speed20")
        src = bce.gap_fill_flags(t, at, aq)
        changed = np.any(np.abs(q_lin - q_sp) > 1e-9, axis=1)
        self.assertTrue(changed.any())
        self.assertTrue(np.all(src[changed] != bce.SRC_MEASURED))

    def test_no_gap_no_flag(self):
        at = np.round(np.arange(0.0, 3.01, 0.1), 3)
        aq = np.cumsum(np.ones((len(at), 6)), axis=0)
        self.assertEqual(int(bce.gap_fill_flags(np.arange(0.0, 3.0, 0.1), at, aq).sum()), 0)

    def test_compat_ignores_state_source(self):
        """お手本にない state_source のキーがあっても check_compat はキーの違いとして数えない."""
        with tempfile.TemporaryDirectory() as d:
            import cv2
            for name, extra in (("ours", True), ("ref", False)):
                ep = Path(d) / name / "e0"
                (ep / "frames").mkdir(parents=True)
                cv2.imwrite(str(ep / "frames" / "000000.jpg"), np.zeros((4, 6, 3), np.uint8))
                kw = {"state_source": np.zeros(1, np.int8)} if extra else {}
                np.savez_compressed(ep / "episode.npz", t=np.zeros(1), t_orig=np.zeros(1),
                                    state=np.zeros((1, 7), np.float32), actions=np.zeros((1, 7), np.float32), **kw)
                (ep / "meta.json").write_text(json.dumps({"fps": 10, "image_size": [6, 4], "state_dim": 7,
                                                          "state_names": [], "task": "x"}))
            self.assertEqual(bce.check_compat(Path(d) / "ours", Path(d) / "ref"), [])


def _old_module():
    """直す前の build_correction_episodes.py (git の HEAD) を別の名前で読み込む."""
    src = subprocess.run(["git", "show", "HEAD:pi05/build_correction_episodes.py"], cwd=_HERE,
                         capture_output=True, text=True, check=True).stdout
    tmp = Path(tempfile.mkdtemp()) / "bce_old.py"
    tmp.write_text(src)
    spec = importlib.util.spec_from_file_location("bce_old", tmp)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


# 補正の多い本・巻き戻しのある本など. 依頼 196 の時点で KIOXIA にあるもの
SAMPLE_RUNS = ("pi05_20261031_175527", "pi05_20261031_180435")


@unittest.skipUnless(Path("/Volumes/KIOXIA").is_mount() and bce.RUNS_ROOT.is_dir(), "KIOXIA がない")
class TestSameAsBefore(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.old = _old_module()

    def _build(self, mod, run, d, gap_fill, trim_hold):
        args = SimpleNamespace(gap_fill=gap_fill, trim_hold=trim_hold, no_trim_mid_idle=False,
                               max_protrude_px=None, force=True)
        out = Path(d) / mod.__name__ / f"{gap_fill}_{trim_hold}" / run
        res = mod.build_one(bce.RUNS_ROOT / run, out, args)
        self.assertTrue(res.get("used"), res)
        z = np.load(out / "episode.npz")
        return {k: z[k] for k in z.files}, json.loads((out / "meta.json").read_text())

    def test_bytes_and_counts(self):
        with tempfile.TemporaryDirectory() as d:
            for run in SAMPLE_RUNS:
                if not (bce.RUNS_ROOT / run).is_dir():
                    self.skipTest(f"{run} がない")
                for gap_fill, trim_hold in (("linear", False), ("speed20", False), ("speed20", True)):
                    with self.subTest(run=run, gap_fill=gap_fill, trim_hold=trim_hold):
                        a, _ = self._build(self.old, run, d, gap_fill, trim_hold)
                        b, meta = self._build(bce, run, d, gap_fill, trim_hold)
                        self.assertNotIn("state_source", a)
                        for k in ("t", "t_orig", "state", "actions"):
                            self.assertEqual(a[k].dtype, b[k].dtype, k)
                            self.assertEqual(a[k].tobytes(), b[k].tobytes(), k)
                        src = b["state_source"]
                        self.assertEqual(src.dtype, np.int8)
                        self.assertEqual(len(src), len(b["state"]))
                        self.assertTrue(set(np.unique(src)) <= set(bce.STATE_SOURCE_CODES))
                        gf = meta["gap_fill"]
                        filled = int((src != 0).sum())
                        self.assertEqual(meta["filled_steps"], filled)
                        self.assertEqual(meta["filled_motion_steps"], int((src == 1).sum()))
                        self.assertAlmostEqual(meta["filled_ratio"], filled / len(src), places=4)
                        # 印のつながりの数は抜けの数を超えない (1 つの抜けが印を 2 つに分けることはない)
                        self.assertLessEqual(runs_of_nonzero(src), gf["gaps"])
                        self.assertLessEqual(runs_of_nonzero(src == 1), gf["motion_gaps"])
                        if gf["gaps"] == 0:
                            self.assertEqual(filled, 0)


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