#!/usr/bin/env python3
"""test_merge_and_config.py — 依頼 45 で足した部分の確かめ (KIOXIA にも Pod にも触らない. 一時フォルダだけ).

  python3 pi05/test_merge_and_config.py

中身:
  - merge_raw_episodes: 2 つの raw_episodes を写してまとめる / 元は変えない / ._* を写さない・tar に入れない /
    形が合わない・名前が重なる・出力先がある ときは止まる
  - ._*.jpg (exFAT の AppleDouble) があっても check_compat が画像の枚数を数え違えない
  - _rmtree: 本体と一緒に ._ が消えても止まらない
  - add_trim_corr_config: trim2 の設定を写してデータ・ステップ数・初期の重み・学習率の予定だけ替える / 二重に入れない /
    違う引数での入れ直しは止まる / 前の版 (corr_v1・corr_v2) が入っていたら止まる
  - pod_prep_trim_corr.sh: 正規化の値を計算し直さず trim2 の値を置く (依頼 48)
"""
import ast
import json
import subprocess
import sys
import tarfile
import tempfile
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))
sys.path.insert(0, str(_HERE / "openpi_addon"))

import add_trim_corr_config as atc  # noqa: E402
import build_correction_episodes as bce  # noqa: E402
import merge_raw_episodes as mre  # noqa: E402


def write_ep(d, n=5, source="teach", size=(384, 216), appledouble=False):
    (d / "frames").mkdir(parents=True)
    for i in range(n):
        cv2.imwrite(str(d / "frames" / f"{i:06d}.jpg"), np.full((size[1], size[0], 3), 90, np.uint8))
        if appledouble:
            (d / "frames" / f"._{i:06d}.jpg").write_bytes(b"\x00\x05\x16\x07")
    st = np.zeros((n, 7), np.float32)
    np.savez_compressed(d / "episode.npz", t=np.arange(n) / 10.0, t_orig=np.arange(n) / 10.0,
                        state=st, actions=st)
    meta = {"task": "put the red block in the white box", "fps": 10, "steps": n, "image_size": list(size),
            "state_dim": 7, "state_names": ["j0", "j1", "j2", "j3", "j4", "j5", "gripper"], "source": source}
    if source.startswith("pi05_policy"):
        meta.update(correct_steps=2, model_steps=n - 2)
    (d / "meta.json").write_text(json.dumps(meta))
    if appledouble:
        (d / "._episode.npz").write_bytes(b"\x00\x05\x16\x07")


def run_merge(root, *extra):
    argv = ["merge_raw_episodes.py", "--data-root", str(root), "--src", "a", "--src", "b",
            "--dataset", "ab", *extra]
    with mock.patch.object(sys, "argv", argv):
        return mre.main()


class TestMerge(unittest.TestCase):
    def setUp(self):
        self._td = tempfile.TemporaryDirectory()
        self.root = Path(self._td.name)
        ds = self.root / "datasets"
        write_ep(ds / "a" / "raw_episodes" / "run_1", n=5)
        write_ep(ds / "a" / "raw_episodes" / "run_2", n=6)
        write_ep(ds / "b" / "raw_episodes" / "pi05_1", n=4, source="pi05_policy (model + correct)",
                 appledouble=True)

    def tearDown(self):
        self._td.cleanup()

    def test_merge_copies_and_counts(self):
        before = sorted(p.relative_to(self.root) for p in (self.root / "datasets").rglob("*"))
        self.assertEqual(run_merge(self.root, "--tar"), 0)
        out = self.root / "datasets" / "ab"
        self.assertEqual(sorted(p.name for p in (out / "raw_episodes").iterdir()), ["pi05_1", "run_1", "run_2"])
        self.assertEqual(list((out / "raw_episodes").rglob("._*")), [])          # AppleDouble は写さない
        idx = json.loads((out / "index.json").read_text())
        self.assertEqual((idx["total_episodes"], idx["total_steps"]), (3, 15))
        self.assertEqual(idx["sources"]["b"]["correct_steps"], 2)
        self.assertEqual(idx["repo_id"], "local/ab")
        self.assertTrue((out / "code" / "merge_raw_episodes.py").exists())        # データと対のコード
        with tarfile.open(out / "ab_raw_episodes.tar") as t:
            names = t.getnames()
        self.assertFalse([n for n in names if Path(n).name.startswith("._")])
        self.assertEqual(sum(n.endswith(".jpg") for n in names), 15)
        # 元は変わらない (増えたのは ab/ の下だけ)
        after = sorted(p.relative_to(self.root) for p in (self.root / "datasets").rglob("*")
                       if not str(p.relative_to(self.root / "datasets")).startswith("ab"))
        self.assertEqual(before, after)

    def test_refuses_existing_and_duplicates(self):
        self.assertEqual(run_merge(self.root), 0)
        with self.assertRaises(SystemExit):
            run_merge(self.root)                                                  # もうある
        self.assertEqual(run_merge(self.root, "--force"), 0)
        write_ep(self.root / "datasets" / "b" / "raw_episodes" / "run_1", n=3)    # 名前が重なる
        with self.assertRaises(SystemExit):
            run_merge(self.root, "--force")

    def test_refuses_mismatch_without_writing(self):
        write_ep(self.root / "datasets" / "b" / "raw_episodes" / "pi05_2", n=3, size=(320, 180))
        with mock.patch("sys.stdout"):
            self.assertEqual(run_merge(self.root), 1)
        self.assertFalse((self.root / "datasets" / "ab").exists())

    def test_check_compat_ignores_appledouble(self):
        ds = self.root / "datasets"
        self.assertEqual(bce.check_compat(ds / "b" / "raw_episodes", ds / "a" / "raw_episodes"), [])

    def test_rmtree_tolerates_vanished_files(self):
        d = self.root / "x"
        (d / "sub").mkdir(parents=True)
        (d / "sub" / "f").write_text("1")
        real_unlink = __import__("os").unlink

        def unlink(path, *a, **k):     # 1 回目で ._ 相当の 2 つ目も消えたことにする
            real_unlink(path, *a, **k)
            raise FileNotFoundError(path)
        with mock.patch("os.unlink", unlink):
            bce._rmtree(d)
        self.assertFalse(d.exists())


# install_addon.py の TRAIN を元に add_trim_config.py (lab/Log11/pod) と同じ置き換えをした trim2 の設定
def fake_config(root):
    src = (_HERE / "openpi_addon" / "install_addon.py").read_text()
    train = src[src.index("TRAIN = f'''") + len("TRAIN = f'''"):]
    train = train[:train.index("'''")].replace("{MARK}", "# >>> r2r2r mecharm addon")
    trim = (train.replace("# >>> r2r2r mecharm addon (train config)", "# >>> r2r2r mecharm teach_v1_trim2 (train config)")
                 .replace("# <<< r2r2r mecharm addon (train config)", "# <<< r2r2r mecharm teach_v1_trim2 (train config)")
                 .replace('name="pi05_mecharm_lora"', 'name="pi05_mecharm_trim_lora"')
                 .replace('repo_id="local/mecharm_redblock_box_v1"',
                          'repo_id="local/mecharm_redblock_box_teach_v1_trim2"')
                 .replace("num_train_steps=5_000,", "num_train_steps=4_000,"))
    cfg = root / "src" / "openpi" / "training" / "config.py"
    cfg.parent.mkdir(parents=True)
    cfg.write_text("_CONFIGS = [\n" + train + trim + "]\n")
    return cfg


def configs(cfg):
    """config.py の TrainConfig(name=...) ごとの キーワード引数 (文字列のまま)."""
    tree = ast.parse(cfg.read_text())
    out = {}
    for node in ast.walk(tree):
        if isinstance(node, ast.Call) and getattr(node.func, "id", "") == "TrainConfig":
            kw = {k.arg: ast.unparse(k.value) for k in node.keywords}
            out[ast.literal_eval(kw["name"])] = kw
    return out


class TestConfig(unittest.TestCase):
    def test_base_and_resume(self):
        for init, loader in (("base", "gs://openpi-assets/checkpoints/pi05_base/params"),
                             ("/workspace/ckpt/3999/params", "/workspace/ckpt/3999/params")):
            with tempfile.TemporaryDirectory() as td, mock.patch("sys.stdout"):
                cfg = fake_config(Path(td))
                self.assertEqual(atc.main([td, "--steps", "1500", "--init", init]), 0)
                c = configs(cfg)
                self.assertEqual(sorted(c), ["pi05_mecharm_lora", "pi05_mecharm_trim_corr_lora",
                                             "pi05_mecharm_trim_lora"])
                new, trim = c["pi05_mecharm_trim_corr_lora"], c["pi05_mecharm_trim_lora"]
                self.assertIn("local/mecharm_redblock_box_teach_v1_trim2_corr_v3", new["data"])   # 依頼 48
                self.assertEqual(new["num_train_steps"], "1500")
                self.assertIn(loader, new["weight_loader"])
                # 依頼 48: 学習率の予定は続きから用に弱める. 下げきる先は trim2 と同じ
                lr = {k.arg: ast.literal_eval(k.value) for k in ast.parse(new["lr_schedule"]).body[0].value.keywords}
                lr_trim = {k.arg: ast.literal_eval(k.value)
                           for k in ast.parse(trim["lr_schedule"]).body[0].value.keywords}
                self.assertEqual(lr, {"warmup_steps": 100, "peak_lr": 2.5e-5, "decay_steps": 1500,
                                      "decay_lr": lr_trim["decay_lr"]})
                self.assertEqual(lr_trim, {"warmup_steps": 300, "peak_lr": 5e-5, "decay_steps": 5000,
                                           "decay_lr": 5e-6})                  # trim2 の設定は変えない
                # ほかは trim2 と同じ
                for k in set(trim) - {"name", "data", "num_train_steps", "weight_loader", "lr_schedule"}:
                    self.assertEqual(new[k], trim[k], k)
                self.assertEqual(trim["num_train_steps"], "4000")              # trim2 の設定は変えない
                text = cfg.read_text()
                self.assertEqual(atc.main([td, "--steps", "1500", "--init", init]), 0)   # 二重に入れない
                self.assertEqual(cfg.read_text(), text)
                with self.assertRaises(SystemExit):
                    atc.main([td, "--steps", "4000", "--init", init])          # 違う引数では止まる
                # データ名・学習率の予定が違うものは止まる (codex r2 推奨: 5 つとも. 止まったときは書かない)
                for x, y in (("_trim2_corr_v3\"", "_trim2_corr_vX\""), ("warmup_steps=100,", "warmup_steps=300,"),
                             ("peak_lr=2.5e-5,", "peak_lr=5e-5,"), ("decay_steps=1_500,", "decay_steps=5_000,"),
                             ("decay_lr=5e-6,\n        ),\n        num_train_steps=1500",
                              "decay_lr=1e-6,\n        ),\n        num_train_steps=1500")):
                    bad = text.replace(x, y, 1) if x.startswith("decay_lr") else \
                        text[:text.index("pi05_mecharm_trim_corr_lora")] + \
                        text[text.index("pi05_mecharm_trim_corr_lora"):].replace(x, y, 1)
                    self.assertNotEqual(bad, text, x)
                    cfg.write_text(bad)
                    with self.assertRaises(SystemExit):
                        atc.main([td, "--steps", "1500", "--init", init])
                    self.assertEqual(cfg.read_text(), bad)

    def test_refuses_old_corr_blocks(self):
        """依頼 46 / 48: 前の版 (corr_v1・corr_v2) の設定が入っていたら足さずに止まる (同じ名前が 2 つにならない)."""
        for tag in ("teach_v1_trim2_corr_v1", "teach_v1_trim2_corr_v2"):
            with tempfile.TemporaryDirectory() as td, mock.patch("sys.stdout"):
                cfg = fake_config(Path(td))
                s = cfg.read_text()
                cfg.write_text(s.replace("]\n", f"    # >>> r2r2r mecharm {tag} (train config)\n"
                                                f"    # <<< r2r2r mecharm {tag} (train config)\n]\n"))
                before = cfg.read_text()
                with self.assertRaises(SystemExit):
                    atc.main([td, "--steps", "1500", "--init", "/workspace/ckpt/3999/params"])
                self.assertEqual(cfg.read_text(), before)

    def test_requires_trim_first(self):
        with tempfile.TemporaryDirectory() as td:
            cfg = Path(td) / "src" / "openpi" / "training" / "config.py"
            cfg.parent.mkdir(parents=True)
            cfg.write_text("_CONFIGS = []\n")
            with self.assertRaises(AssertionError):
                atc.main([td, "--steps", "10", "--init", "base"])

    def test_cli_requires_steps_and_init(self):
        r = subprocess.run([sys.executable, str(_HERE / "openpi_addon" / "add_trim_corr_config.py"), "/tmp"],
                           capture_output=True, text=True)
        self.assertNotEqual(r.returncode, 0)


class TestPodPrep(unittest.TestCase):
    """依頼 48: pod_prep_trim_corr.sh は正規化の値を計算し直さず, 3999 の学習の trim2 の値を置く."""
    SH = _HERE / "openpi_addon" / "pod_prep_trim_corr.sh"
    MAC_NORM = Path("/Volumes/KIOXIA/r2r2r-data/checkpoints/pi05_mecharm_trim_lora_v1/3999/assets/local/"
                    "mecharm_redblock_box_teach_v1_trim2/norm_stats.json")

    def lines(self):
        return [ln for ln in self.SH.read_text().splitlines() if not ln.lstrip().startswith("#")]

    def test_no_compute_norm_stats(self):
        self.assertFalse(any("compute_norm_stats" in ln for ln in self.lines()))

    def test_places_trim2_norm(self):
        code = "\n".join(self.lines())
        self.assertIn("NAME=mecharm_redblock_box_teach_v1_trim2_corr_v3", code)
        self.assertIn("NORM_SRC=/workspace/ckpt/3999/assets/local/mecharm_redblock_box_teach_v1_trim2/norm_stats.json",
                      code)
        self.assertIn("NORM_DST=assets/$CONFIG/local/$NAME/norm_stats.json", code)   # openpi が学習で読む場所
        self.assertIn('cp "$NORM_SRC" "$NORM_DST"', code)

    def test_sha_matches_mac_copy(self):
        import hashlib
        import re
        m = re.search(r"^NORM_SHA256=([0-9a-f]{64})", self.SH.read_text(), re.M)
        self.assertIsNotNone(m)
        if not self.MAC_NORM.exists():
            self.skipTest("KIOXIA がない")
        self.assertEqual(hashlib.sha256(self.MAC_NORM.read_bytes()).hexdigest(), m.group(1))

    def test_bash_syntax(self):
        self.assertEqual(subprocess.run(["bash", "-n", str(self.SH)]).returncode, 0)


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