#!/usr/bin/env python3

"""
Copyright (c) 2025 Mohamed Aboelnasr
Licensed under the GNU Affero General Public License v3.0 or later (AGPL-3.0-or-later)
"""
import unittest
import numpy as np
from unittest.mock import patch, MagicMock
from ManipulaPy.kinematics import SerialManipulator
from ManipulaPy import utils
import tempfile
import os


class TestKinematics(unittest.TestCase):
    """
    Comprehensive tests for the kinematics module in ManipulaPy.
    Tests cover forward kinematics, inverse kinematics, Jacobian calculations,
    end-effector velocity, joint velocity, and error handling.
    """

    def setUp(self):
        """Set up test fixtures with a standard 6-DOF robot configuration."""
        # 1) Screw axes in space frame (6,6)
        self.Slist = np.array(
            [
                [0, 0, 1, 0, 0, 0],
                [0, -1, 0, -0.089, 0, 0],
                [0, -1, 0, -0.089, 0, 0.425],
                [0, -1, 0, -0.089, 0, 0.817],
                [1, 0, 0, 0, 0.109, 0],
                [0, -1, 0, -0.089, 0, 0.817],
            ]
        ).T  # shape => (6,6)

        # 2) Home configuration (M)
        self.M = np.array(
            [[1, 0, 0, 0.817], [0, 1, 0, 0], [0, 0, 1, 0.191], [0, 0, 0, 1]]
        )

        # 3) Omega list from top 3 rows (for the constructor)
        self.omega_list = self.Slist[:3, :]

        # 4) Body frame screw axes B_list
        self.B_list = np.copy(self.Slist)

        # 5) Joint limits for testing
        self.joint_limits = [(-np.pi, np.pi)] * 6

        # 6) Create the SerialManipulator
        self.robot = SerialManipulator(
            M_list=self.M,
            omega_list=self.omega_list,
            S_list=self.Slist,
            B_list=self.B_list,
            joint_limits=self.joint_limits,
        )

        # 7) Test joint angles
        self.test_angles = np.array([0.1, 0.2, -0.3, 0.4, -0.5, 0.6])
        self.zero_angles = np.zeros(6)

    def test_constructor_with_all_parameters(self):
        """Test constructor with all parameters specified."""
        r_list = np.random.rand(3, 6)
        b_list = np.random.rand(3, 6)
        G_list = np.random.rand(6, 6, 6)
        
        robot = SerialManipulator(
            M_list=self.M,
            omega_list=self.omega_list,
            r_list=r_list,
            b_list=b_list,
            S_list=self.Slist,
            B_list=self.B_list,
            G_list=G_list,
            joint_limits=self.joint_limits,
        )
        
        np.testing.assert_array_equal(robot.r_list, r_list)
        np.testing.assert_array_equal(robot.b_list, b_list)
        np.testing.assert_array_equal(robot.G_list, G_list)
        self.assertEqual(robot.joint_limits, self.joint_limits)

    def test_constructor_with_minimal_parameters(self):
        """Test constructor with minimal parameters (automatic extraction)."""
        # Create proper r_list and b_list for this test
        r_list = np.array([[0, 0, 0, 0, 0.109, 0],
                          [0, -0.089, -0.089, -0.089, 0, -0.089],
                          [0, 0, 0.425, 0.817, 0, 0.817]]).T  # (6, 3) -> transpose to (3, 6)
        
        robot = SerialManipulator(
            M_list=self.M,
            omega_list=self.omega_list,
            S_list=self.Slist,
            r_list=r_list,  # Provide r_list to avoid extraction issues
        )
        
        # Should have proper r_list
        self.assertIsNotNone(robot.r_list)
        # Should have created default joint limits
        self.assertEqual(len(robot.joint_limits), 6)
        self.assertEqual(robot.joint_limits[0], (None, None))

    def test_constructor_without_s_and_b_lists(self):
        """Test constructor automatic generation of S_list and B_list."""
        # Provide proper r_list and b_list to avoid extraction from empty arrays
        r_list = np.zeros((3, 6))  # 3x6 array of zeros
        b_list = np.zeros((3, 6))  # 3x6 array of zeros
        
        robot = SerialManipulator(
            M_list=self.M,
            omega_list=self.omega_list,
            r_list=r_list,
            b_list=b_list,
        )
        
        # Should have generated S_list and B_list
        self.assertIsNotNone(robot.S_list)
        self.assertIsNotNone(robot.B_list)

    def test_update_state(self):
        """Test the update_state method."""
        joint_positions = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
        joint_velocities = np.array([0.01, 0.02, 0.03, 0.04, 0.05, 0.06])
        
        # Test with both positions and velocities
        self.robot.update_state(joint_positions, joint_velocities)
        np.testing.assert_array_equal(self.robot.joint_positions, joint_positions)
        np.testing.assert_array_equal(self.robot.joint_velocities, joint_velocities)
        
        # Test with only positions (velocities should be zero)
        new_positions = np.array([0.2, 0.3, 0.4, 0.5, 0.6, 0.7])
        self.robot.update_state(new_positions)
        np.testing.assert_array_equal(self.robot.joint_positions, new_positions)
        np.testing.assert_array_equal(self.robot.joint_velocities, np.zeros(6))

    def test_forward_kinematics_space(self):
        """Test forward kinematics in space frame."""
        # Zero angles should give M
        T_space = self.robot.forward_kinematics(self.zero_angles, frame="space")
        np.testing.assert_array_almost_equal(
            T_space,
            self.M,
            decimal=4,
            err_msg="Space frame at zero angles should match M.",
        )
        
        # Non-zero angles should give different result
        T_space_nonzero = self.robot.forward_kinematics(self.test_angles, frame="space")
        self.assertFalse(np.allclose(T_space_nonzero, self.M))
        
        # Result should be a valid homogeneous transformation matrix
        self.assertEqual(T_space_nonzero.shape, (4, 4))
        # Last row should be [0, 0, 0, 1]
        np.testing.assert_array_almost_equal(T_space_nonzero[3, :], [0, 0, 0, 1])
        # Rotation part should be orthogonal
        R = T_space_nonzero[:3, :3]
        np.testing.assert_array_almost_equal(R @ R.T, np.eye(3), decimal=4)

    def test_forward_kinematics_body(self):
        """Test forward kinematics in body frame."""
        # Zero angles should give M
        T_body = self.robot.forward_kinematics(self.zero_angles, frame="body")
        np.testing.assert_array_almost_equal(
            T_body,
            self.M,
            decimal=4,
            err_msg="Body frame at zero angles should match M.",
        )
        
        # Non-zero angles
        T_body_nonzero = self.robot.forward_kinematics(self.test_angles, frame="body")
        self.assertEqual(T_body_nonzero.shape, (4, 4))

    def test_forward_kinematics_invalid_frame(self):
        """Test forward kinematics with invalid frame."""
        with self.assertRaises(ValueError) as context:
            self.robot.forward_kinematics(self.test_angles, frame="invalid")
        self.assertIn("Invalid frame specified", str(context.exception))

    def test_jacobian_space_frame(self):
        """Test Jacobian calculation in space frame."""
        J = self.robot.jacobian(self.test_angles, frame="space")
        
        # Should be 6 x n_joints
        self.assertEqual(J.shape, (6, 6))
        
        # Jacobian should not be all zeros for non-zero joint angles
        self.assertFalse(np.allclose(J, np.zeros((6, 6))))

    def test_jacobian_body_frame(self):
        """Test Jacobian calculation in body frame."""
        J = self.robot.jacobian(self.test_angles, frame="body")
        
        # Should be 6 x n_joints
        self.assertEqual(J.shape, (6, 6))

    def test_jacobian_invalid_frame(self):
        """Test Jacobian with invalid frame."""
        with self.assertRaises(ValueError) as context:
            self.robot.jacobian(self.test_angles, frame="invalid")
        self.assertIn("Invalid frame specified", str(context.exception))

    def test_end_effector_velocity_space(self):
        """Test end-effector velocity calculation in space frame."""
        dthetalist = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
        
        V_ee = self.robot.end_effector_velocity(self.test_angles, dthetalist, frame="space")
        
        # Should be a 6-element vector (3 angular + 3 linear velocity)
        self.assertEqual(V_ee.shape, (6,))

    def test_end_effector_velocity_body(self):
        """Test end-effector velocity calculation in body frame."""
        dthetalist = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
        
        V_ee = self.robot.end_effector_velocity(self.test_angles, dthetalist, frame="body")
        
        # Should be a 6-element vector
        self.assertEqual(V_ee.shape, (6,))

    def test_end_effector_velocity_invalid_frame(self):
        """Test end-effector velocity with invalid frame."""
        dthetalist = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
        
        with self.assertRaises(ValueError) as context:
            self.robot.end_effector_velocity(self.test_angles, dthetalist, frame="invalid")
        self.assertIn("Invalid frame specified", str(context.exception))

    def test_joint_velocity_space(self):
        """Test joint velocity calculation in space frame."""
        V_ee = np.array([0.1, 0.2, 0.3, 0.01, 0.02, 0.03])
        
        joint_vel = self.robot.joint_velocity(self.test_angles, V_ee, frame="space")
        
        # Should be same length as number of joints
        self.assertEqual(joint_vel.shape, (6,))

    def test_joint_velocity_body(self):
        """Test joint velocity calculation in body frame."""
        V_ee = np.array([0.1, 0.2, 0.3, 0.01, 0.02, 0.03])
        
        joint_vel = self.robot.joint_velocity(self.test_angles, V_ee, frame="body")
        
        # Should be same length as number of joints
        self.assertEqual(joint_vel.shape, (6,))

    def test_joint_velocity_invalid_frame(self):
        """Test joint velocity with invalid frame."""
        V_ee = np.array([0.1, 0.2, 0.3, 0.01, 0.02, 0.03])
        
        with self.assertRaises(ValueError) as context:
            self.robot.joint_velocity(self.test_angles, V_ee, frame="invalid")
        self.assertIn("Invalid frame specified", str(context.exception))

    def test_end_effector_pose(self):
        """Test end-effector pose calculation."""
        pose = self.robot.end_effector_pose(self.test_angles)
        
        # Should be 6-element vector [x, y, z, roll, pitch, yaw]
        self.assertEqual(pose.shape, (6,))
        
        # Position should be the first 3 elements
        T = self.robot.forward_kinematics(self.test_angles)
        expected_position = T[:3, 3]
        np.testing.assert_array_almost_equal(pose[:3], expected_position)

    def test_simple_inverse_kinematics(self):
        """Test simple inverse kinematics case."""
        target_pose = np.copy(self.M)
        init_guess = np.zeros(6)

        solution, success, iterations = self.robot.iterative_inverse_kinematics(
            T_desired=target_pose,
            thetalist0=init_guess,
            eomg=1e-5,
            ev=1e-5,
            max_iterations=500,
        )
        
        self.assertTrue(success, "IK solver did not converge to a solution for M.")
        self.assertIsInstance(iterations, int)
        self.assertGreater(iterations, 0)

        final_pose = self.robot.forward_kinematics(solution, frame="space")
        np.testing.assert_array_almost_equal(
            final_pose,
            target_pose,
            decimal=3,
            err_msg="IK solution's forward kinematics does not match target.",
        )

    def test_inverse_kinematics_with_joint_limits(self):
        """Test inverse kinematics with joint limits enforcement."""
        # Set tight joint limits
        robot_limited = SerialManipulator(
            M_list=self.M,
            omega_list=self.omega_list,
            S_list=self.Slist,
            B_list=self.B_list,
            joint_limits=[(-0.1, 0.1)] * 6,  # Very tight limits
        )
        
        # Create target that requires large joint movement
        target_pose = np.copy(self.M)
        target_pose[0, 3] += 0.5  # Move in X direction
        
        init_guess = np.zeros(6)
        solution, success, _ = robot_limited.iterative_inverse_kinematics(
            T_desired=target_pose,
            thetalist0=init_guess,
            max_iterations=100,
        )
        
        # Check that joint limits are respected
        for i, (min_limit, max_limit) in enumerate(robot_limited.joint_limits):
            if min_limit is not None:
                self.assertGreaterEqual(solution[i], min_limit)
            if max_limit is not None:
                self.assertLessEqual(solution[i], max_limit)

    def test_inverse_kinematics_convergence_failure(self):
        """Test inverse kinematics when convergence fails."""
        # Create impossible target (very far away)
        target_pose = np.copy(self.M)
        target_pose[:3, 3] = [100, 100, 100]  # Unreachable position
        
        init_guess = np.zeros(6)
        solution, success, iterations = self.robot.iterative_inverse_kinematics(
            T_desired=target_pose,
            thetalist0=init_guess,
            max_iterations=10,  # Very low iteration limit
        )
        
        # Should fail to converge
        self.assertFalse(success)
        self.assertEqual(iterations, 11)  # max_iterations + 1

    def test_inverse_kinematics_step_capping(self):
        """Test inverse kinematics with step size capping."""
        target_pose = np.copy(self.M)
        target_pose[:3, 3] = [0.5, 0.5, 0.5]  # Moderate displacement
        
        init_guess = np.ones(6) * 0.1  # Non-zero starting guess
        solution, success, _ = self.robot.iterative_inverse_kinematics(
            T_desired=target_pose,
            thetalist0=init_guess,
            step_cap=0.1,  # Small step cap
            max_iterations=1000,
        )
        
        # Should still converge with small steps
        if success:  # Only check if convergence was achieved
            final_pose = self.robot.forward_kinematics(solution)
            position_error = np.linalg.norm(final_pose[:3, 3] - target_pose[:3, 3])
            self.assertLess(position_error, 0.1)

    @patch('matplotlib.pyplot.savefig')
    @patch('matplotlib.pyplot.close')
    @patch('matplotlib.pyplot.tight_layout')
    @patch('matplotlib.pyplot.grid')
    @patch('matplotlib.pyplot.legend')
    @patch('matplotlib.pyplot.title')
    @patch('matplotlib.pyplot.ylabel')
    @patch('matplotlib.pyplot.xlabel')
    @patch('matplotlib.pyplot.plot')
    @patch('matplotlib.use')
    def test_inverse_kinematics_with_plot(self, mock_use, mock_plot, mock_xlabel, 
                                         mock_ylabel, mock_title, mock_legend, 
                                         mock_grid, mock_tight_layout, mock_close, 
                                         mock_savefig):
        """Test inverse kinematics with residual plotting."""
        target_pose = np.copy(self.M)
        init_guess = np.zeros(6)
        
        with tempfile.TemporaryDirectory() as temp_dir:
            png_path = os.path.join(temp_dir, "test_residuals.png")
            
            solution, success, _ = self.robot.iterative_inverse_kinematics(
                T_desired=target_pose,
                thetalist0=init_guess,
                plot_residuals=True,
                png_name=png_path,
                max_iterations=50,
            )
            
            # Verify matplotlib functions were called
            mock_use.assert_called_with("Agg")
            mock_plot.assert_called()
            mock_savefig.assert_called_with(png_path, dpi=400)
            mock_close.assert_called()

    def test_inverse_kinematics_damping_parameter(self):
        """Test inverse kinematics with different damping values."""
        target_pose = np.copy(self.M)
        target_pose[0, 3] += 0.1  # Small displacement
        
        init_guess = np.zeros(6)
        
        # Test with high damping
        solution_high, success_high, iter_high = self.robot.iterative_inverse_kinematics(
            T_desired=target_pose,
            thetalist0=init_guess,
            damping=0.5,  # High damping
            max_iterations=1000,
        )
        
        # Test with low damping
        solution_low, success_low, iter_low = self.robot.iterative_inverse_kinematics(
            T_desired=target_pose,
            thetalist0=init_guess,
            damping=1e-6,  # Low damping
            max_iterations=1000,
        )
        
        # Both should converge but may take different iterations
        if success_high and success_low:
            # Verify both solutions are valid
            pose_high = self.robot.forward_kinematics(solution_high)
            pose_low = self.robot.forward_kinematics(solution_low)
            
            error_high = np.linalg.norm(pose_high[:3, 3] - target_pose[:3, 3])
            error_low = np.linalg.norm(pose_low[:3, 3] - target_pose[:3, 3])
            
            self.assertLess(error_high, 0.01)
            self.assertLess(error_low, 0.01)

    def test_velocity_consistency(self):
        """Test consistency between forward kinematics and velocity calculations."""
        # Use small time step for numerical differentiation
        dt = 1e-7  # Even smaller time step for better accuracy
        dthetalist = np.array([0.01, 0.02, 0.03, 0.04, 0.05, 0.06])  # Smaller joint velocities
        
        # Current and slightly perturbed joint angles
        theta1 = self.test_angles
        theta2 = theta1 + dthetalist * dt
        
        # Forward kinematics at both configurations
        T1 = self.robot.forward_kinematics(theta1)
        T2 = self.robot.forward_kinematics(theta2)
        
        # Numerical differentiation of position
        p1 = T1[:3, 3]
        p2 = T2[:3, 3]
        numerical_velocity = (p2 - p1) / dt
        
        # Analytical velocity from Jacobian
        J = self.robot.jacobian(theta1)
        
        # Check the Jacobian structure first
        self.assertEqual(J.shape, (6, 6), "Jacobian should be 6x6")
        
        # Linear velocity is in rows 3:6 (bottom 3 rows)
        analytical_velocity = J[3:6, :] @ dthetalist  # Linear part only
        
        # Check if the magnitudes are reasonable
        analytical_mag = np.linalg.norm(analytical_velocity)
        numerical_mag = np.linalg.norm(numerical_velocity)
        
        # If magnitudes are vastly different, there might be a frame issue
        if abs(analytical_mag - numerical_mag) > max(analytical_mag, numerical_mag):
            # Try with body frame Jacobian
            J_body = self.robot.jacobian(theta1, frame="body")
            analytical_velocity_body = J_body[3:6, :] @ dthetalist
            
            # Test which one is closer
            error_space = np.linalg.norm(analytical_velocity - numerical_velocity)
            error_body = np.linalg.norm(analytical_velocity_body - numerical_velocity)
            
            if error_body < error_space:
                analytical_velocity = analytical_velocity_body
        
        # Use relative tolerance for better robustness
        max_error = max(np.linalg.norm(analytical_velocity), np.linalg.norm(numerical_velocity))
        relative_tolerance = 0.1  # 10% relative tolerance
        
        try:
            # First try with decimal=1
            np.testing.assert_array_almost_equal(
                analytical_velocity, numerical_velocity, decimal=1
            )
        except AssertionError:
            # If that fails, check if the relative error is acceptable
            absolute_error = np.linalg.norm(analytical_velocity - numerical_velocity)
            relative_error = absolute_error / max_error if max_error > 1e-10 else absolute_error
            
            if relative_error > relative_tolerance:
                # Print debug information
                print(f"\nVelocity consistency test debug info:")
                print(f"Analytical velocity: {analytical_velocity}")
                print(f"Numerical velocity:  {numerical_velocity}")
                print(f"Absolute error: {absolute_error}")
                print(f"Relative error: {relative_error}")
                print(f"Joint angles: {theta1}")
                print(f"Joint velocities: {dthetalist}")
                print(f"Time step: {dt}")
                
                # Try even more relaxed tolerance
                np.testing.assert_array_almost_equal(
                    analytical_velocity, numerical_velocity, decimal=0
                )


    def test_jacobian_consistency_between_frames(self):
        """Test that Jacobian calculations are consistent between frames."""
        J_space = self.robot.jacobian(self.test_angles, frame="space")
        J_body = self.robot.jacobian(self.test_angles, frame="body")
        
        # Both should have same shape
        self.assertEqual(J_space.shape, J_body.shape)
        
        # They should generally be different (unless at zero configuration)
        if not np.allclose(self.test_angles, 0):
            self.assertFalse(np.allclose(J_space, J_body))

    def test_edge_cases_with_small_angles(self):
        """Test edge cases with very small joint angles."""
        small_angles = np.array([1e-10, -1e-10, 1e-12, -1e-12, 1e-8, -1e-8])
        
        # Should not crash and should be close to zero configuration
        T = self.robot.forward_kinematics(small_angles)
        np.testing.assert_array_almost_equal(T, self.M, decimal=6)

    def test_edge_cases_with_large_angles(self):
        """Test edge cases with large joint angles."""
        large_angles = np.array([10*np.pi, -10*np.pi, 5*np.pi, -5*np.pi, 20*np.pi, -20*np.pi])
        
        # Should not crash
        T = self.robot.forward_kinematics(large_angles)
        self.assertEqual(T.shape, (4, 4))
        
        # Jacobian should also work
        J = self.robot.jacobian(large_angles)
        self.assertEqual(J.shape, (6, 6))

    def test_consistency_across_methods(self):
        """Test consistency between different methods."""
        # Test that forward kinematics and end_effector_pose are consistent
        T = self.robot.forward_kinematics(self.test_angles)
        pose = self.robot.end_effector_pose(self.test_angles)
        
        # Position should match
        np.testing.assert_array_almost_equal(T[:3, 3], pose[:3])

    def tearDown(self):
        """Clean up after tests."""
        # Clean up any temporary files or resources if needed
        pass


if __name__ == "__main__":
    # Run tests with verbose output
    unittest.main(verbosity=2)