"""Optional installed-mjlab checks; GPU smoke is opt-in, never a full training run."""

import importlib.util
import os
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace

SERVICE_ROOT = Path(__file__).resolve().parents[1]
for root in (SERVICE_ROOT, SERVICE_ROOT / "rl"):
    sys.path.insert(0, str(root))

HAS_MJLAB = importlib.util.find_spec("mjlab") is not None


@unittest.skipUnless(HAS_MJLAB, "mjlab is not installed in this Python")
class ObstacleContractTest(unittest.TestCase):
    def test_training_config_file_is_revalidated(self):
        import json
        import tempfile

        from scripts.train import _load_task_config
        from task_config import OBSTACLE_TASK, TaskConfigError, validate_task_config

        with tempfile.TemporaryDirectory() as directory:
            source = Path(directory) / "training_config.json"
            custom = validate_task_config(OBSTACLE_TASK, {"sensorCfg": {"fov": 60}}, 42)
            source.write_text(json.dumps(custom))
            self.assertEqual(_load_task_config(OBSTACLE_TASK, str(source), 42), custom)
            with self.assertRaises(ValueError):
                _load_task_config(OBSTACLE_TASK, str(source), 43)
            custom["sensorCfg"]["maxDistance"] = -1
            source.write_text(json.dumps(custom))
            with self.assertRaises(TaskConfigError):
                _load_task_config(OBSTACLE_TASK, str(source), 42)
            source.write_text("x" * (128 * 1024 + 1))
            with self.assertRaises(ValueError):
                _load_task_config(OBSTACLE_TASK, str(source), 42)

    def test_pattern_order_normalization_and_body_offsets(self):
        import torch
        from src.tasks.obstacle_avoidance.mdp import ForwardFanPatternCfg, forward_depth

        offsets, directions = ForwardFanPatternCfg(fov=90).generate_rays(None, "cpu")
        self.assertEqual(tuple(directions.shape), (32, 3))
        torch.testing.assert_close(offsets, torch.tensor([0.3, 0, 0.05]).repeat(32, 1))
        torch.testing.assert_close(directions.norm(dim=1), torch.ones(32))
        self.assertAlmostEqual(directions[0, 1].item(), -(0.5**0.5), places=6)
        self.assertAlmostEqual(directions[-1, 1].item(), 0.5**0.5, places=6)
        scene = {
            "forward_scan": SimpleNamespace(
                data=SimpleNamespace(
                    distances=torch.tensor([[-1, 0, 1, 4, 8]], dtype=torch.float32)
                )
            )
        }
        result = forward_depth(SimpleNamespace(scene=scene), max_distance=4)
        torch.testing.assert_close(result, torch.tensor([[1, 0, 0.25, 1, 1]]))

    def test_navigation_matches_body_heading_and_arrival_stop(self):
        import torch
        from src.tasks.obstacle_avoidance.mdp import NavigationCommandCfg

        data = SimpleNamespace(
            root_link_pos_w=torch.tensor([[-5.0, 0, 0.32]]),
            root_link_quat_w=torch.tensor([[1.0, 0, 0, 0]]),
        )

        class Scene(dict):
            env_origins = torch.tensor([[-5.0, 0, 0]])

        env = SimpleNamespace(
            num_envs=1, device="cpu", scene=Scene(robot=SimpleNamespace(data=data))
        )
        command = NavigationCommandCfg(resampling_time_range=(1e9, 1e9)).build(env)
        torch.testing.assert_close(command.command, torch.tensor([[0.6, 0, 0]]))
        self.assertAlmostEqual(command.errors()[1].item(), 10)
        data.root_link_quat_w[:] = torch.tensor([[0.5**0.5, 0, 0, 0.5**0.5]])
        self.assertAlmostEqual(command.command[0, 2].item(), -1)
        data.root_link_pos_w[0, 0] = 4.8
        torch.testing.assert_close(command.command, torch.zeros((1, 3)))

    def test_navigation_reset_randomizes_safe_distant_pairs_and_heading(self):
        import torch
        from src.tasks.obstacle_avoidance.mdp import NavigationCommandCfg

        class Robot:
            def __init__(self, count):
                self.data = SimpleNamespace(
                    default_root_state=torch.tensor([[0, 0, 0.32] + [0] * 10] * count),
                    root_link_pos_w=torch.zeros((count, 3)),
                    root_link_quat_w=torch.zeros((count, 4)),
                )

            def write_root_link_pose_to_sim(self, pose, env_ids):
                self.data.root_link_pos_w[env_ids] = pose[:, :3]
                self.data.root_link_quat_w[env_ids] = pose[:, 3:]

            def write_root_link_velocity_to_sim(self, velocity, env_ids):
                self.velocity = velocity

        class Scene(dict):
            env_origins = torch.zeros((64, 3))

        robot = Robot(64)
        env = SimpleNamespace(num_envs=64, device="cpu", scene=Scene(robot=robot))
        command = NavigationCommandCfg(
            resampling_time_range=(1e9, 1e9),
            navigation_points=((-2, 0), (-1, 0), (0, 0), (1, 0), (2, 0)),
            component_starts=(0,),
            component_counts=(5,),
            fallback_pairs=(((-2, 0), (2, 0)),),
            min_goal_distance=2,
        ).build(env)
        torch.manual_seed(7)
        command.sample_episode(torch.arange(64))
        self.assertTrue(
            ((command.goals_w - robot.data.root_link_pos_w[:, :2]).norm(dim=1) >= 2).all()
        )
        self.assertGreater(torch.unique(robot.data.root_link_pos_w[:, :2], dim=0).shape[0], 1)
        self.assertGreater(torch.unique(command.goals_w, dim=0).shape[0], 1)
        torch.testing.assert_close(robot.data.root_link_quat_w.norm(dim=1), torch.ones(64))
        self.assertTrue((robot.velocity == 0).all())

    def test_all_terrain_presets_compile_at_exported_world_coordinates(self):
        import mujoco
        import numpy as np
        from mjlab.terrains import TerrainGenerator
        from src.tasks.obstacle_avoidance.env_cfg import unitree_go2_obstacle_env_cfg
        from src.tasks.obstacle_avoidance.terrain import apply_terrain_configuration
        from task_config import (
            OBSTACLE_TASK,
            TERRAIN_PRESETS,
            build_terrain_layout,
            validate_task_config,
        )

        for preset in (p for p in TERRAIN_PRESETS if p != "custom_boxes"):
            with self.subTest(preset=preset):
                custom = validate_task_config(OBSTACLE_TASK, {"terrainPreset": preset}, 7)
                cfg = unitree_go2_obstacle_env_cfg()
                apply_terrain_configuration(cfg, custom)
                generator = TerrainGenerator(cfg.scene.terrain.terrain_generator)
                spec = mujoco.MjSpec()
                generator.compile(spec)
                model = spec.compile()
                layout = build_terrain_layout(custom)
                np.testing.assert_allclose(model.geom_pos, [box["pos"] for box in layout["boxes"]])
                np.testing.assert_allclose(
                    model.geom_size, [box["size"] for box in layout["boxes"]]
                )
                np.testing.assert_allclose(generator.terrain_origins[0, 0], [-5, 0, 0])
                np.testing.assert_allclose(model.geom_friction[:, 0], layout["friction"])
                self.assertTrue((model.geom_group == 0).all())

    @unittest.skipUnless(os.environ.get("GO2_RUN_MJLAB_SMOKE") == "1", "opt-in GPU smoke")
    def test_real_environment_81_observation_and_mujoco_raycast_parity(self):
        import mujoco
        import numpy as np
        import torch
        import warp as wp
        from mjlab.envs import ManagerBasedRlEnv
        from src.tasks.obstacle_avoidance.env_cfg import (
            apply_obstacle_configuration,
            unitree_go2_obstacle_env_cfg,
        )
        from src.tasks.obstacle_avoidance.mdp import ForwardFanPatternCfg
        from task_config import JOINT_NAMES, OBSTACLE_TASK, validate_task_config
        from warp._src import context

        if not torch.cuda.is_available():
            self.skipTest("CUDA is unavailable")
        if not hasattr(wp, "context"):
            wp.context = context
        cfg = unitree_go2_obstacle_env_cfg()
        apply_obstacle_configuration(
            cfg, validate_task_config(OBSTACLE_TASK, {}, 42), randomize_navigation=False
        )
        cfg.scene.num_envs = 2
        env = ManagerBasedRlEnv(cfg, device="cuda:0")
        try:
            obs, _ = env.reset()
            self.assertEqual(tuple(obs["actor"].shape), (2, 81))
            self.assertEqual(list(env.scene["robot"].joint_names), JOINT_NAMES)
            np.testing.assert_allclose(env.scene.env_origins.cpu(), [[-5, 0, 0]] * 2)
            np.testing.assert_allclose(
                env.scene["robot"].data.root_link_pos_w.cpu(), [[-5, 0, 0.32]] * 2, atol=1e-6
            )
            for _ in range(3):
                obs, reward, terminated, truncated, _ = env.step(
                    torch.zeros((2, 12), device=env.device)
                )
                self.assertTrue(torch.isfinite(obs["actor"]).all())
                self.assertTrue(torch.isfinite(reward).all())
                self.assertFalse(terminated.any() or truncated.any())
            model = env.sim.mj_model
            data = mujoco.MjData(model)
            data.qpos[:] = env.sim.data.qpos[0].cpu().numpy()
            mujoco.mj_forward(model, data)
            body = model.body("robot/base_link").id
            rotation = data.xmat[body].reshape(3, 3)
            offsets, directions = ForwardFanPatternCfg().generate_rays(None, "cpu")
            expected = []
            for offset, direction in zip(offsets.numpy(), directions.numpy(), strict=True):
                distance = mujoco.mj_ray(
                    model,
                    data,
                    data.xpos[body] + rotation @ offset,
                    rotation @ direction,
                    np.array([1, 0, 0, 0, 0, 0], dtype=np.uint8),
                    1,
                    -1,
                    np.array([-1], dtype=np.int32),
                )
                expected.append(1 if distance < 0 else min(1, distance / 4))
            np.testing.assert_allclose(obs["actor"][0, 47:79].cpu(), expected, atol=2e-5)
        finally:
            env.close()


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