import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock

import numpy as np

from roboimi.utils import raw_action_trajectory_viewer as traj_view


class RawActionTrajectoryViewerTest(unittest.TestCase):
    def test_set_transfer_box_pose_writes_joint_qpos(self):
        joint_qpos = np.zeros(7, dtype=np.float64)

        class _FakeJoint:
            def __init__(self, qpos):
                self.qpos = qpos

        class _FakeData:
            def joint(self, name):
                assert name == "red_box_joint"
                return _FakeJoint(joint_qpos)

        traj_view.set_transfer_box_pose(_FakeData(), np.array([0.2, -0.1, 1.05], dtype=np.float64))

        np.testing.assert_array_equal(
            joint_qpos,
            np.array([0.2, -0.1, 1.05, 1.0, 0.0, 0.0, 0.0], dtype=np.float64),
        )

    def test_disable_cv2_highgui_temporarily_replaces_gui_calls(self):
        fake_cv2 = SimpleNamespace(
            namedWindow=lambda *args, **kwargs: "named",
            imshow=lambda *args, **kwargs: "imshow",
            waitKey=lambda *args, **kwargs: "wait",
        )

        restore = traj_view.disable_cv2_highgui(fake_cv2)
        self.assertIsNone(fake_cv2.namedWindow("x"))
        self.assertIsNone(fake_cv2.imshow("x", None))
        self.assertEqual(fake_cv2.waitKey(1), 1)

        restore()
        self.assertEqual(fake_cv2.namedWindow("x"), "named")
        self.assertEqual(fake_cv2.imshow("x", None), "imshow")
        self.assertEqual(fake_cv2.waitKey(1), "wait")

    def test_load_raw_action_positions_from_npz(self):
        raw_action = np.array(
            [
                [1.0, 2.0, 3.0, 0, 0, 0, 1, 11.0, 12.0, 13.0, 0, 0, 0, 1, -1, -1],
                [4.0, 5.0, 6.0, 0, 0, 0, 1, 14.0, 15.0, 16.0, 0, 0, 0, 1, -1, -1],
            ],
            dtype=np.float32,
        )
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "trajectory.npz"
            np.savez(path, raw_action=raw_action)

            positions = traj_view.load_raw_action_positions(path)

        np.testing.assert_array_equal(
            positions["left"],
            np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32),
        )
        np.testing.assert_array_equal(
            positions["right"],
            np.array([[11.0, 12.0, 13.0], [14.0, 15.0, 16.0]], dtype=np.float32),
        )

    def test_build_red_capsule_segments_downsamples_to_fit_scene_limit(self):
        left = np.stack([np.array([float(i), 0.0, 0.0], dtype=np.float32) for i in range(6)])
        right = np.stack([np.array([float(i), 1.0, 0.0], dtype=np.float32) for i in range(6)])

        markers = traj_view.build_trajectory_capsule_markers(
            {"left": left, "right": right},
            max_markers=4,
            radius=0.01,
        )

        self.assertLessEqual(len(markers), 4)
        self.assertTrue(all(marker["rgba"] == (1.0, 0.0, 0.0, 1.0) for marker in markers))
        self.assertTrue(all(marker["radius"] == 0.01 for marker in markers))

    def test_apply_capsule_markers_populates_user_scene(self):
        fake_scene = SimpleNamespace(
            maxgeom=3,
            ngeom=99,
            geoms=[object(), object(), object()],
        )
        markers = [
            {
                "from": np.array([0.0, 0.0, 0.0], dtype=np.float64),
                "to": np.array([1.0, 0.0, 0.0], dtype=np.float64),
                "rgba": (1.0, 0.0, 0.0, 1.0),
                "radius": 0.01,
            },
            {
                "from": np.array([0.0, 1.0, 0.0], dtype=np.float64),
                "to": np.array([1.0, 1.0, 0.0], dtype=np.float64),
                "rgba": (1.0, 0.0, 0.0, 1.0),
                "radius": 0.01,
            },
        ]

        with mock.patch.object(traj_view.mujoco, "mjv_initGeom") as init_geom, mock.patch.object(
            traj_view.mujoco,
            "mjv_connector",
        ) as connector:
            traj_view.apply_capsule_markers_to_scene(fake_scene, markers)

        self.assertEqual(fake_scene.ngeom, 2)
        self.assertEqual(init_geom.call_count, 2)
        self.assertEqual(connector.call_count, 2)


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