import unittest

import mujoco
import numpy as np
from robot_descriptions.loaders.mujoco import load_robot_description

import mjpl


class TestMinkIKSolver(unittest.TestCase):
    def test_ik(self):
        model = load_robot_description("ur5e_mj_description")
        site_name = "attachment_site"
        constraints = [
            mjpl.JointLimitConstraint(model),
            mjpl.CollisionConstraint(model),
        ]

        q_init = model.keyframe("home").qpos.copy()

        # Generate a pose for the IK problem by running FK on a
        # random configuration within joint limits.
        data = mujoco.MjData(model)
        rng = np.random.default_rng(seed=12345)
        data.qpos = rng.uniform(*model.jnt_range.T)
        mujoco.mj_kinematics(model, data)
        target_pose = mjpl.site_pose(data, site_name)

        pos_tolerance = 1e-3
        ori_tolerance = 1e-3
        solver = mjpl.MinkIKSolver(
            model=model,
            joints=mjpl.all_joints(model),
            constraints=constraints,
            pos_tolerance=pos_tolerance,
            ori_tolerance=ori_tolerance,
            seed=12345,
            max_attempts=5,
        )

        # Solve IK (test both with/without an initial guess config)
        ik_solutions = []
        ik_solutions.extend(
            solver.solve_ik(
                pose=target_pose,
                site=site_name,
                q_init_guess=q_init,
            )
        )
        ik_solutions.extend(
            solver.solve_ik(
                pose=target_pose,
                site=site_name,
                q_init_guess=None,
            )
        )
        self.assertEqual(len(ik_solutions), 2)
        self.assertTrue(mjpl.obeys_constraints(ik_solutions[0], constraints))
        self.assertTrue(mjpl.obeys_constraints(ik_solutions[1], constraints))

        # Confirm that the IK solution gives a site pose within the specified error.
        # We cannot rely on checking similarity between q_candidate and the randomly
        # generated config since there may be multiple solutions to the IK problem!
        for solution in ik_solutions:
            data.qpos = solution.copy()
            mujoco.mj_kinematics(model, data)
            actual_site_pose = mjpl.site_pose(data, site_name)
            err = target_pose.minus(actual_site_pose)
            self.assertLessEqual(np.linalg.norm(err[:3]), pos_tolerance)
            self.assertLessEqual(np.linalg.norm(err[3:]), ori_tolerance)

    def test_ik_subset_joints(self):
        model = load_robot_description("ur5e_mj_description")
        site_name = "attachment_site"
        constraints = [
            mjpl.JointLimitConstraint(model),
            mjpl.CollisionConstraint(model),
        ]

        q_init = model.keyframe("home").qpos.copy()

        # Generate a pose for the IK problem that's derived from a random config
        # that was generated by modifying the subset of joints used for solving IK.
        modifiable_joints = ["shoulder_pan_joint", "elbow_joint"]
        q_rand = mjpl.random_config(
            model, q_init, modifiable_joints, seed=12345, constraints=constraints
        )
        data = mujoco.MjData(model)
        data.qpos = q_rand
        mujoco.mj_kinematics(model, data)
        target_pose = mjpl.site_pose(data, site_name)

        # Solve IK for the target pose, using only a subset of joints.
        solver = mjpl.MinkIKSolver(
            model=model,
            joints=modifiable_joints,
            constraints=constraints,
            seed=12345,
            max_attempts=5,
        )
        solutions = solver.solve_ik(target_pose, site_name, q_init)
        self.assertEqual(len(solutions), 1)
        self.assertTrue(mjpl.obeys_constraints(solutions[0], constraints))

        # Check that only the specified subset of joints were modified in order to solve IK.
        fixed_idx = [
            i
            for i in range(model.nq)
            if i not in mjpl.qpos_idx(model, modifiable_joints)
        ]
        np.testing.assert_allclose(
            solutions[0][fixed_idx], q_init[fixed_idx], rtol=0, atol=1e-12
        )

    def test_invalid_args(self):
        model = load_robot_description("ur5e_mj_description")
        joints = mjpl.all_joints(model)

        with self.assertRaisesRegex(ValueError, "`max_attempts` must be > 0"):
            mjpl.MinkIKSolver(
                model=model,
                joints=joints,
                max_attempts=-2,
            )
            mjpl.MinkIKSolver(
                model=model,
                joints=joints,
                max_attempts=0,
            )

        with self.assertRaisesRegex(ValueError, "`iterations` must be > 0"):
            mjpl.MinkIKSolver(
                model=model,
                joints=joints,
                iterations=-2,
            )
            mjpl.MinkIKSolver(
                model=model,
                joints=joints,
                iterations=0,
            )

        with self.assertRaisesRegex(ValueError, "cannot be empty"):
            mjpl.MinkIKSolver(
                model=model,
                joints=[],
            )


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