"""CPU transfer regression; real source and <=4env/1iteration GPU checks are opt-in."""

import copy
import importlib.util
import io
import json
import os
import sys
import tempfile
import unittest
from dataclasses import asdict, replace
from pathlib import Path
from types import SimpleNamespace

ROOT = Path(__file__).resolve().parents[1]
for root in (ROOT, ROOT / "rl"):
    sys.path.insert(0, str(root))
HAS_STACK = all(importlib.util.find_spec(m) is not None for m in ("mjlab", "torch", "onnxruntime"))


@unittest.skipUnless(HAS_STACK, "installed RL + CPU ORT stack required")
class PretrainedTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        import torch

        torch.set_num_threads(1)
        from pretrained import ValidatedSource, make_reference_actor
        from src.tasks.velocity.config.go2.env_cfgs import unitree_go2_flat_env_cfg
        from src.tasks.velocity.config.go2.rl_cfg import unitree_go2_ppo_runner_cfg

        torch.manual_seed(42)
        actor = make_reference_actor()
        with torch.no_grad():
            actor.obs_normalizer._mean.uniform_(-0.2, 0.2)
            actor.obs_normalizer._var.uniform_(0.1, 1.2)
            actor.obs_normalizer._std.copy_(actor.obs_normalizer._var.sqrt())
            actor.obs_normalizer.count.fill_(983138304)
        cls.source = ValidatedSource(actor.state_dict(), {"source_id": "synthetic"})
        cls.env_cfg = asdict(unitree_go2_flat_env_cfg())
        cls.agent_cfg = asdict(unitree_go2_ppo_runner_cfg())

    def test_extension_preserves_actions_and_learns_new_columns(self):
        import torch
        from pretrained import comparison_observations, make_reference_actor, warm_start_actor
        from tensordict import TensorDict

        source_actor = make_reference_actor().eval()
        source_actor.load_state_dict(self.source.actor_state)
        base = comparison_observations()
        expected = source_actor.mlp(source_actor.obs_normalizer(base)).detach()
        for dim in (47, 81, 97):
            with self.subTest(dim=dim):
                actor = make_reference_actor(dim)
                warm_start_actor(actor, self.source)
                for key in ("_mean", "_var", "_std"):
                    value = getattr(actor.obs_normalizer, key)
                    self.assertTrue(
                        torch.equal(value[:, :47], self.source.actor_state[f"obs_normalizer.{key}"])
                    )
                    self.assertTrue(
                        torch.equal(
                            value[:, 47:],
                            torch.full_like(value[:, 47:], 0 if key == "_mean" else 1),
                        )
                    )
                self.assertTrue(
                    torch.equal(
                        actor.obs_normalizer.count, self.source.actor_state["obs_normalizer.count"]
                    )
                )
                self.assertEqual(actor.mlp[0].weight[:, 47:].count_nonzero().item(), 0)
                self.assertTrue(actor.mlp[0].weight.requires_grad)
                extra = torch.rand(len(base), dim - 47) * 2 - 1
                x = torch.cat((base, extra), dim=1)
                torch.testing.assert_close(
                    actor.mlp(actor.obs_normalizer(x)), expected, atol=3e-6, rtol=3e-6
                )
                initial_error = (actor.mlp(actor.obs_normalizer(x)) - expected).abs().max().item()
                x[:, 47:] *= -1
                torch.testing.assert_close(
                    actor.mlp(actor.obs_normalizer(x)), expected, atol=3e-6, rtol=3e-6
                )
                before = actor.obs_normalizer._mean.clone()
                obs = TensorDict({"actor": x}, batch_size=[len(base)])
                actor.update_normalization(obs)
                self.assertEqual(actor.obs_normalizer.count.item(), 983138304 + len(base))
                self.assertLess(
                    (actor.obs_normalizer._mean[:, :47] - before[:, :47]).abs().max().item(), 1e-6
                )
                # Updating statistics is not eval identity. Random gravity probes include
                # out-of-distribution values in nearly constant source channels.
                updated_source = copy.deepcopy(source_actor).train()
                updated_source.update_normalization(
                    TensorDict({"actor": base}, batch_size=[len(base)])
                )
                updated_expected = updated_source.mlp(updated_source.obs_normalizer(base)).detach()
                torch.testing.assert_close(actor(obs), updated_expected, atol=3e-6, rtol=3e-6)
                self.assertLess((actor(obs) - expected).abs().max().item(), 2e-4)
                if hasattr(self, "evidence"):
                    self.evidence.setdefault("cpu_transfer", {})[str(dim)] = {
                        "initial_max_abs_error": initial_error,
                        "updated_source_max_abs_error": (actor(obs) - updated_expected)
                        .abs()
                        .max()
                        .item(),
                        "random_batch_action_max_abs_change": (actor(obs) - expected)
                        .abs()
                        .max()
                        .item(),
                    }
                loss = actor(obs).square().mean()
                loss.backward()
                grad = actor.mlp[0].weight.grad
                self.assertTrue(torch.isfinite(grad).all())
                if dim > 47:
                    self.assertGreater(grad[:, 47:].abs().max().item(), 0)
                torch.optim.Adam(actor.parameters(), lr=1e-4).step()
                if dim > 47:
                    self.assertGreater(actor.mlp[0].weight[:, 47:].abs().max().item(), 0)
                stream = io.BytesIO()
                torch.save(actor.state_dict(), stream)
                stream.seek(0)
                restored = make_reference_actor(dim)
                restored.load_state_dict(torch.load(stream, weights_only=True), strict=True)
                for k, v in actor.state_dict().items():
                    self.assertTrue(torch.equal(v, restored.state_dict()[k]), k)

    def test_rejects_shapes_nonfinite_and_runtime_activation(self):
        import torch
        from pretrained import (
            PretrainedError,
            ValidatedSource,
            make_reference_actor,
            warm_start_actor,
        )

        for key, value in (
            ("mlp.0.weight", torch.zeros(512, 46)),
            ("obs_normalizer._std", torch.full((1, 47), float("nan"))),
            ("distribution.std_param", torch.zeros(12)),
        ):
            state = dict(self.source.actor_state)
            state[key] = value
            with self.assertRaises(PretrainedError):
                warm_start_actor(make_reference_actor(81), ValidatedSource(state, {}))
        actor = make_reference_actor(81)
        actor.mlp[1] = torch.nn.ReLU()
        with self.assertRaisesRegex(PretrainedError, "architecture"):
            warm_start_actor(actor, self.source)
        with self.assertRaises(PretrainedError):
            warm_start_actor(make_reference_actor(82), self.source)

    def test_semantics_fail_closed(self):
        from pretrained import PretrainedError, _plain, validate_semantics

        env, agent = _plain(self.env_cfg), _plain(self.agent_cfg)
        validate_semantics(env, agent, self.env_cfg, self.agent_cfg)
        for edit in (
            lambda e: e["observations"]["actor"]["terms"]["phase"]["params"].update(period=0.7),
            lambda e: e["observations"]["actor"]["terms"]["joint_vel"].update(scale=0.1),
            lambda e: e["actions"]["joint_pos"].update(scale=0.5),
            lambda e: e["scene"]["entities"]["robot"]["articulation"]["actuators"][0].update(
                armature=0.1
            ),
        ):
            bad = copy.deepcopy(env)
            edit(bad)
            with self.assertRaises(PretrainedError):
                validate_semantics(bad, agent, self.env_cfg, self.agent_cfg)
            with self.assertRaises(PretrainedError):
                validate_semantics(env, agent, bad, self.agent_cfg)
        bad_agent = copy.deepcopy(agent)
        bad_agent["actor"]["activation"] = "relu"
        with self.assertRaises(PretrainedError):
            validate_semantics(env, bad_agent, self.env_cfg, self.agent_cfg)

    def test_compiled_runtime_contract(self):
        from unittest.mock import patch

        from pretrained import BASE_TERMS, JOINTS, PretrainedError, validate_runtime_contract

        metadata = {
            "joint_names": JOINTS,
            "observation_names": BASE_TERMS + ["forward_depth", "target_error"],
            "command_names": ["twist"],
            "action_scale": 0.25,
            "joint_stiffness": [20, 20, 40] * 4,
            "joint_damping": [1, 1, 2] * 4,
            "default_joint_pos": [-0.1, 0.9, -1.8, 0.1, 0.9, -1.8] * 2,
        }
        with patch("mjlab.rl.exporter_utils.get_base_metadata", return_value=metadata):
            validate_runtime_contract(None)
            metadata["joint_names"] = list(reversed(JOINTS))
            with self.assertRaisesRegex(PretrainedError, "joint order"):
                validate_runtime_contract(None)

    def test_safe_yaml_and_allowed_roots(self):
        from pretrained import PretrainedError, _read_allowed, _yaml_data

        with self.assertRaises(PretrainedError):
            _yaml_data(b"x: !!python/object/apply:os.system ['touch /tmp/never-pretrained']")
        with self.assertRaises(PretrainedError):
            _yaml_data(b"x: 1\nx: 2")
        self.assertEqual(_yaml_data(b"axis: {0: 1, 1: 2}"), {"axis": {0: 1, 1: 2}})
        self.assertEqual(
            _yaml_data(b"x: !!python/name:os.system ''"), {"x": {"symbol": "os.system"}}
        )
        with tempfile.TemporaryDirectory() as directory, tempfile.TemporaryDirectory() as outside:
            root = Path(directory)
            secret = Path(outside) / "secret.pt"
            secret.write_bytes(b"x")
            (root / "link.pt").symlink_to(secret)
            for path in (secret, root / "link.pt", root / ".." / Path(outside).name / "secret.pt"):
                with self.assertRaises(PretrainedError):
                    _read_allowed(path, [root], 128)
            with self.assertRaises(PretrainedError):
                _read_allowed(secret, [Path(outside)], 0)

    def test_fresh_runner_only_and_trial_consistency(self):
        import torch
        from pretrained import PretrainedError, initialize_runner, make_reference_actor

        actors = []
        for _ in range(2):
            actor = make_reference_actor(81)
            critic = torch.nn.Linear(108, 1)
            original = copy.deepcopy(critic.state_dict())
            optimizer = torch.optim.Adam(list(actor.parameters()) + list(critic.parameters()))
            runner = SimpleNamespace(
                current_learning_iteration=0,
                alg=SimpleNamespace(actor=actor, critic=critic, optimizer=optimizer),
            )
            initialize_runner(runner, self.source)
            self.assertFalse(optimizer.state)
            self.assertEqual(runner.current_learning_iteration, 0)
            for k in original:
                self.assertTrue(torch.equal(original[k], critic.state_dict()[k]))
            actors.append(actor.state_dict())
            runner.current_learning_iteration = 1
            with self.assertRaises(PretrainedError):
                initialize_runner(runner, self.source)
            runner.current_learning_iteration = 0
            optimizer.state[actor.mlp[0].weight] = {"step": torch.tensor(1)}
            with self.assertRaises(PretrainedError):
                initialize_runner(runner, self.source)
        for k in actors[0]:
            self.assertTrue(torch.equal(actors[0][k], actors[1][k]))

    def test_cli_modes_do_not_reinitialize_resume(self):
        from scripts.train import TrainConfig, _load_pretrained
        from src.tasks.velocity.config.go2.env_cfgs import unitree_go2_flat_env_cfg
        from src.tasks.velocity.config.go2.rl_cfg import unitree_go2_ppo_runner_cfg

        cfg = TrainConfig(env=unitree_go2_flat_env_cfg(), agent=unitree_go2_ppo_runner_cfg())
        self.assertIsNone(_load_pretrained(cfg))
        self.assertIsNone(_load_pretrained(replace(cfg, resume_checkpoint="trial.pt")))
        with self.assertRaisesRegex(ValueError, "mutually exclusive"):
            _load_pretrained(
                replace(cfg, resume_checkpoint="trial.pt", pretrained_checkpoint="source.pt")
            )
        cfg.agent.resume = True
        with self.assertRaisesRegex(ValueError, "mutually exclusive"):
            _load_pretrained(replace(cfg, pretrained_checkpoint="source.pt"))
        with self.assertRaisesRegex(ValueError, "ONNX alone"):
            _load_pretrained(replace(cfg, pretrained_onnx="policy.onnx"))


@unittest.skipUnless(
    HAS_STACK and os.environ.get("GO2_PRETRAINED_SOURCE"),
    "set GO2_PRETRAINED_SOURCE to explicit allowed source directory",
)
class RealPretrainedTest(PretrainedTest):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        from pretrained import read_pretrained_source

        directory = Path(os.environ["GO2_PRETRAINED_SOURCE"])
        cls.source = read_pretrained_source(
            directory / "model_10000.pt",
            allowed_roots=[directory],
            target_env=cls.env_cfg,
            target_agent=cls.agent_cfg,
        )
        cls.evidence = {"source": cls.source.manifest}

    @classmethod
    def tearDownClass(cls):
        if os.environ.get("GO2_PRETRAINED_EVIDENCE_DIR"):
            directory = Path(os.environ["GO2_PRETRAINED_EVIDENCE_DIR"])
            directory.mkdir(parents=True, exist_ok=True)
            (directory / "real-source-evidence.json").write_text(
                json.dumps(cls.evidence, indent=2) + "\n"
            )

    def test_wrong_checkpoint_actor_rejected_by_onnx(self):
        from pretrained import PretrainedError, make_reference_actor, verify_onnx

        # No bulk loading: a fresh random actor is sufficient to prove fail-closed identity.
        onnx = (Path(os.environ["GO2_PRETRAINED_SOURCE"]) / "policy.onnx").read_bytes()
        with self.assertRaisesRegex(PretrainedError, "does not match"):
            verify_onnx(onnx, make_reference_actor())

    def test_extended_export_matches_cpu_actor(self):
        import numpy as np
        import onnxruntime as ort
        import torch
        from pretrained import comparison_observations, make_reference_actor, warm_start_actor

        errors = {}
        with tempfile.TemporaryDirectory() as directory:
            for dim in (81, 97):
                actor = make_reference_actor(dim).eval()
                warm_start_actor(actor, self.source)
                export = actor.as_onnx(verbose=False)
                path = str(Path(directory) / f"policy-{dim}.onnx")
                torch.onnx.export(
                    export,
                    export.get_dummy_inputs(),
                    path,
                    input_names=export.input_names,
                    output_names=export.output_names,
                    opset_version=18,
                    dynamo=False,
                )
                session = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
                x = torch.cat((comparison_observations(), torch.rand(48, dim - 47) * 2 - 1), dim=1)
                with torch.no_grad():
                    expected = export(x).numpy()
                actual = np.concatenate(
                    [session.run(None, {"obs": row[None].numpy()})[0] for row in x]
                )
                np.testing.assert_allclose(actual, expected, atol=2e-5, rtol=2e-5)
                errors[str(dim)] = float(np.abs(actual - expected).max())
        self.evidence["extended_onnx_max_abs_error"] = errors

    @unittest.skipUnless(
        os.environ.get("GO2_PRETRAINED_GPU_SMOKE") == "1", "opt-in <=4env x 1iteration GPU smoke"
    )
    def test_real_observations_and_one_ppo_iteration(self):
        import torch
        import warp as wp

        if not hasattr(wp, "context"):
            from warp._src import context

            wp.context = context
        from mjlab.envs import ManagerBasedRlEnv
        from mjlab.rl import RslRlVecEnvWrapper
        from pretrained import initialize_runner, make_reference_actor, validate_runtime_contract
        from src.tasks.obstacle_avoidance.env_cfg import unitree_go2_obstacle_env_cfg
        from src.tasks.velocity.config.go2.rl_cfg import unitree_go2_ppo_runner_cfg
        from src.tasks.velocity.rl.runner import VelocityOnPolicyRunner
        from task_config import OBSTACLE_TASK, deployment_metadata, validate_task_config
        from tensordict import TensorDict

        env_cfg = unitree_go2_obstacle_env_cfg()
        env_cfg.scene.num_envs = 4
        env_cfg.seed = 42
        agent_cfg = unitree_go2_ppo_runner_cfg()
        agent_cfg.logger = "tensorboard"
        agent_cfg.max_iterations = 1
        with tempfile.TemporaryDirectory() as directory:
            raw = ManagerBasedRlEnv(env_cfg, device="cuda:0")
            self.evidence["compiled_runtime"] = validate_runtime_contract(raw)
            raw.platform_deployment = deployment_metadata(
                OBSTACLE_TASK, validate_task_config(OBSTACLE_TASK, {}, 42), 42
            )
            env = RslRlVecEnvWrapper(raw)
            try:
                runner = VelocityOnPolicyRunner(env, asdict(agent_cfg), directory, "cuda:0")
                manifest = initialize_runner(runner, self.source)
                actor = runner.alg.actor
                cpu_actor = make_reference_actor().eval()
                cpu_actor.load_state_dict(self.source.actor_state)
                obs = env.get_observations()
                base = obs["actor"][:, :47].cpu()
                with torch.no_grad():
                    expected = cpu_actor(TensorDict({"actor": base}, batch_size=[4]))
                    before = actor(obs).cpu()
                torch.testing.assert_close(before, expected, atol=2e-5, rtol=2e-5)
                old_mean = actor.obs_normalizer._mean.clone()
                actor.update_normalization(obs)
                mean_change = (
                    (actor.obs_normalizer._mean[:, :47] - old_mean[:, :47]).abs().max().item()
                )
                after = actor(obs).detach().cpu()
                self.assertTrue(torch.isfinite(after).all())
                torch.testing.assert_close(before, after, atol=2e-5, rtol=2e-5)
                loss = actor(obs).square().mean()
                loss.backward()
                grad = actor.mlp[0].weight.grad[:, 47:]
                self.assertTrue(torch.isfinite(grad).all())
                self.assertGreater(grad.abs().max().item(), 0)
                grad_max = grad.abs().max().item()
                runner.alg.optimizer.zero_grad()
                critic_before = copy.deepcopy(runner.alg.critic.state_dict())
                runner.learn(num_learning_iterations=1, init_at_random_ep_len=True)
                self.assertGreater(actor.mlp[0].weight[:, 47:].abs().max().item(), 0)
                self.assertTrue(runner.alg.optimizer.state)
                self.assertTrue(
                    any(
                        not torch.equal(v, runner.alg.critic.state_dict()[k])
                        for k, v in critic_before.items()
                    )
                )
                saved = torch.load(
                    Path(directory) / "model_0.pt", map_location="cpu", weights_only=True
                )
                # Same-trial resume uses the installed PPO loader: restore all states.
                actor_state = copy.deepcopy(actor.state_dict())
                # Production promotion starts a fresh process/runner; do not reload
                # into tensors created by the previous rollout's inference_mode.
                resumed = VelocityOnPolicyRunner(env, asdict(agent_cfg), directory, "cuda:0")
                self.assertTrue(resumed.alg.load(saved, None, strict=True))
                resumed.current_learning_iteration = saved["iter"] + 1
                self.assertEqual(resumed.current_learning_iteration, 1)
                self.assertTrue(resumed.alg.optimizer.state)
                actor = resumed.alg.actor
                for k, v in actor_state.items():
                    self.assertTrue(torch.equal(v.cpu(), actor.state_dict()[k].cpu()), k)
                import numpy as np
                import onnxruntime as ort

                session = ort.InferenceSession(
                    str(Path(directory) / "policy.onnx"), providers=["CPUExecutionProvider"]
                )
                final_obs = env.get_observations()
                with torch.no_grad():
                    expected_final = actor(final_obs).cpu().numpy()
                actual = np.concatenate(
                    [
                        session.run(None, {"obs": row[None].cpu().numpy()})[0]
                        for row in final_obs["actor"]
                    ]
                )
                np.testing.assert_allclose(actual, expected_final, atol=2e-5, rtol=2e-5)
                # Independently run source ORT on observations actually sampled from MuJoCo.
                source_session = ort.InferenceSession(
                    str(Path(os.environ["GO2_PRETRAINED_SOURCE"]) / "policy.onnx"),
                    providers=["CPUExecutionProvider"],
                )
                source_actions = np.concatenate(
                    [source_session.run(None, {"obs": row[None].numpy()})[0] for row in base]
                )
                np.testing.assert_allclose(source_actions, before.numpy(), atol=2e-5, rtol=2e-5)
                self.evidence["gpu_smoke"] = {
                    "num_envs": 4,
                    "iterations": 1,
                    "initialization": manifest,
                    "real_observation_source_onnx_max_abs_error": float(
                        np.abs(source_actions - before.numpy()).max()
                    ),
                    "normalization_batch_action_max_abs_change": (after - before)
                    .abs()
                    .max()
                    .item(),
                    "normalization_batch_base_mean_max_abs_change": mean_change,
                    "new_columns_gradient_max": grad_max,
                    "new_columns_weight_max_after_ppo": actor.mlp[0]
                    .weight[:, 47:]
                    .abs()
                    .max()
                    .item(),
                    "export_max_abs_error": float(np.abs(actual - expected_final).max()),
                    "count_restored": actor.obs_normalizer.count.item(),
                    "all_actor_tensors_restored_exactly": True,
                }
            finally:
                env.close()


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