"""Inverse Kinematics controller for UR5e using damped least squares."""
import numpy as np
import mujoco


class IKController:
    """
    IK controller using MuJoCo's built-in Jacobian computation
    and damped least squares for stable inverse kinematics.
    """

    def __init__(self, model, data, ee_site_name="ee_site", arm_joint_names=None):
        """
        Initialize IK controller.

        Args:
            model: MuJoCo model
            data: MuJoCo data
            ee_site_name: Name of the end-effector site
            arm_joint_names: List of joint names to control
        """
        self.model = model
        self.data = data

        # Get site ID
        self.ee_site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, ee_site_name)
        if self.ee_site_id < 0:
            raise ValueError(f"Site '{ee_site_name}' not found in model")

        # Get joint IDs
        self.joint_ids = []
        self.dof_ids = []

        if arm_joint_names is None:
            arm_joint_names = [
                "shoulder_pan_joint",
                "shoulder_lift_joint",
                "elbow_joint",
                "wrist_1_joint",
                "wrist_2_joint",
                "wrist_3_joint",
            ]

        for name in arm_joint_names:
            jnt_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name)
            if jnt_id < 0:
                raise ValueError(f"Joint '{name}' not found in model")
            self.joint_ids.append(jnt_id)
            # Get the DOF address for this joint
            self.dof_ids.append(model.jnt_dofadr[jnt_id])

        self.num_joints = len(self.joint_ids)

        # IK parameters
        self.damping = 0.01  # Damping factor for stability
        self.max_iterations = 50
        self.tolerance = 0.001  # Position tolerance in meters
        self.step_size = 0.5  # Step size for gradient descent

        # Joint limits (from model)
        self.joint_limits_low = np.array([model.jnt_range[jid, 0] for jid in self.joint_ids])
        self.joint_limits_high = np.array([model.jnt_range[jid, 1] for jid in self.joint_ids])

        # Preallocate Jacobian arrays
        self.jacp = np.zeros((3, model.nv))
        self.jacr = np.zeros((3, model.nv))

    def get_current_joint_positions(self):
        """Get current positions of controlled joints."""
        return np.array([self.data.qpos[self.model.jnt_qposadr[jid]] for jid in self.joint_ids])

    def get_ee_position(self):
        """Get current end-effector position."""
        return self.data.site_xpos[self.ee_site_id].copy()

    def get_ee_orientation(self):
        """Get current end-effector orientation as rotation matrix."""
        return self.data.site_xmat[self.ee_site_id].reshape(3, 3).copy()

    def get_ee_quat(self):
        """Get current end-effector orientation as quaternion [w, x, y, z]."""
        xmat = self.get_ee_orientation()
        quat = np.zeros(4)
        mujoco.mju_mat2Quat(quat, xmat.flatten())
        return quat

    def compute_jacobian(self, include_rotation=False):
        """Compute the Jacobian for the end-effector site.

        Args:
            include_rotation: If True, return full 6x6 Jacobian (position + rotation)

        Returns:
            If include_rotation: 6xN Jacobian (position rows, then rotation rows)
            Otherwise: 3xN position-only Jacobian
        """
        # Zero out arrays
        self.jacp.fill(0)
        self.jacr.fill(0)

        # Compute Jacobian using MuJoCo
        mujoco.mj_jacSite(self.model, self.data, self.jacp, self.jacr, self.ee_site_id)

        if include_rotation:
            # Full 6-DOF Jacobian
            jac = np.zeros((6, self.num_joints))
            for i, dof_id in enumerate(self.dof_ids):
                jac[:3, i] = self.jacp[:, dof_id]  # Position
                jac[3:, i] = self.jacr[:, dof_id]  # Rotation
            return jac
        else:
            # Position-only Jacobian
            jac = np.zeros((3, self.num_joints))
            for i, dof_id in enumerate(self.dof_ids):
                jac[:, i] = self.jacp[:, dof_id]
            return jac

    def compute_ik(self, target_pos, current_qpos=None):
        """
        Compute joint positions to reach target position using damped least squares.

        Args:
            target_pos: Target position [x, y, z]
            current_qpos: Starting joint positions (uses current if None)

        Returns:
            Joint positions to reach target
        """
        if current_qpos is None:
            qpos = self.get_current_joint_positions()
        else:
            qpos = current_qpos.copy()

        target = np.array(target_pos)

        for _ in range(self.max_iterations):
            # Update forward kinematics
            for i, jid in enumerate(self.joint_ids):
                self.data.qpos[self.model.jnt_qposadr[jid]] = qpos[i]
            mujoco.mj_forward(self.model, self.data)

            # Get current end-effector position
            ee_pos = self.get_ee_position()

            # Compute error
            error = target - ee_pos
            error_norm = np.linalg.norm(error)

            # Check convergence
            if error_norm < self.tolerance:
                break

            # Compute Jacobian
            jac = self.compute_jacobian()

            # Damped least squares: dq = J^T * (J * J^T + lambda^2 * I)^-1 * error
            jjt = jac @ jac.T
            damped = jjt + self.damping**2 * np.eye(3)
            dq = jac.T @ np.linalg.solve(damped, error)

            # Apply step
            qpos = qpos + self.step_size * dq

            # Clamp to joint limits
            qpos = np.clip(qpos, self.joint_limits_low, self.joint_limits_high)

        return qpos

    def quat_error(self, quat_current, quat_target):
        """Compute orientation error as axis-angle representation.

        Args:
            quat_current: Current quaternion [w, x, y, z]
            quat_target: Target quaternion [w, x, y, z]

        Returns:
            Orientation error as 3D vector (axis * angle)
        """
        # Compute quaternion difference: q_err = q_target * q_current^(-1)
        # For unit quaternions, inverse is conjugate: [w, -x, -y, -z]
        qc = quat_current
        qt = quat_target

        # Quaternion conjugate of current
        qc_conj = np.array([qc[0], -qc[1], -qc[2], -qc[3]])

        # Quaternion multiplication: q_err = qt * qc_conj
        q_err = np.zeros(4)
        q_err[0] = qt[0]*qc_conj[0] - qt[1]*qc_conj[1] - qt[2]*qc_conj[2] - qt[3]*qc_conj[3]
        q_err[1] = qt[0]*qc_conj[1] + qt[1]*qc_conj[0] + qt[2]*qc_conj[3] - qt[3]*qc_conj[2]
        q_err[2] = qt[0]*qc_conj[2] - qt[1]*qc_conj[3] + qt[2]*qc_conj[0] + qt[3]*qc_conj[1]
        q_err[3] = qt[0]*qc_conj[3] + qt[1]*qc_conj[2] - qt[2]*qc_conj[1] + qt[3]*qc_conj[0]

        # Ensure we take the shortest path (w >= 0)
        if q_err[0] < 0:
            q_err = -q_err

        # Convert to axis-angle: angle = 2 * acos(w), axis = [x,y,z] / sin(angle/2)
        # For small angles, error ≈ 2 * [x, y, z]
        angle = 2.0 * np.arccos(np.clip(q_err[0], -1.0, 1.0))
        if angle < 1e-6:
            return np.zeros(3)

        axis = q_err[1:4] / np.sin(angle / 2.0)
        return axis * angle

    def compute_ik_with_orientation(self, target_pos, target_quat=None, current_qpos=None,
                                     position_weight=1.0, orientation_weight=0.5):
        """
        Compute joint positions to reach target pose (position + orientation).

        Args:
            target_pos: Target position [x, y, z]
            target_quat: Target orientation quaternion [w, x, y, z] (optional)
            current_qpos: Starting joint positions (uses current if None)
            position_weight: Weight for position error (default 1.0)
            orientation_weight: Weight for orientation error (default 0.5)

        Returns:
            Joint positions to reach target
        """
        # If no orientation specified, fall back to position-only IK
        if target_quat is None:
            return self.compute_ik(target_pos, current_qpos)

        if current_qpos is None:
            qpos = self.get_current_joint_positions()
        else:
            qpos = current_qpos.copy()

        target_p = np.array(target_pos)
        target_q = np.array(target_quat)

        # Normalize target quaternion
        target_q = target_q / np.linalg.norm(target_q)

        # Tolerance for orientation (radians)
        orientation_tolerance = 0.01

        for _ in range(self.max_iterations):
            # Update forward kinematics
            for i, jid in enumerate(self.joint_ids):
                self.data.qpos[self.model.jnt_qposadr[jid]] = qpos[i]
            mujoco.mj_forward(self.model, self.data)

            # Get current end-effector pose
            ee_pos = self.get_ee_position()
            ee_quat = self.get_ee_quat()

            # Compute position error
            pos_error = target_p - ee_pos
            pos_error_norm = np.linalg.norm(pos_error)

            # Compute orientation error
            ori_error = self.quat_error(ee_quat, target_q)
            ori_error_norm = np.linalg.norm(ori_error)

            # Check convergence
            if pos_error_norm < self.tolerance and ori_error_norm < orientation_tolerance:
                break

            # Compute full 6-DOF Jacobian
            jac = self.compute_jacobian(include_rotation=True)

            # Construct weighted error vector [position; orientation]
            error = np.concatenate([
                position_weight * pos_error,
                orientation_weight * ori_error
            ])

            # Apply weights to Jacobian rows
            jac_weighted = jac.copy()
            jac_weighted[:3, :] *= position_weight
            jac_weighted[3:, :] *= orientation_weight

            # Damped least squares: dq = J^T * (J * J^T + lambda^2 * I)^-1 * error
            jjt = jac_weighted @ jac_weighted.T
            damped = jjt + self.damping**2 * np.eye(6)
            dq = jac_weighted.T @ np.linalg.solve(damped, error)

            # Apply step
            qpos = qpos + self.step_size * dq

            # Clamp to joint limits
            qpos = np.clip(qpos, self.joint_limits_low, self.joint_limits_high)

        return qpos

    @staticmethod
    def euler_to_quat(roll, pitch, yaw):
        """Convert Euler angles (XYZ convention) to quaternion [w, x, y, z].

        Args:
            roll: Rotation around X axis (radians)
            pitch: Rotation around Y axis (radians)
            yaw: Rotation around Z axis (radians)

        Returns:
            Quaternion [w, x, y, z]
        """
        cy = np.cos(yaw * 0.5)
        sy = np.sin(yaw * 0.5)
        cp = np.cos(pitch * 0.5)
        sp = np.sin(pitch * 0.5)
        cr = np.cos(roll * 0.5)
        sr = np.sin(roll * 0.5)

        w = cr * cp * cy + sr * sp * sy
        x = sr * cp * cy - cr * sp * sy
        y = cr * sp * cy + sr * cp * sy
        z = cr * cp * sy - sr * sp * cy

        return np.array([w, x, y, z])

    @staticmethod
    def quat_to_euler(quat):
        """Convert quaternion [w, x, y, z] to Euler angles (XYZ convention).

        Args:
            quat: Quaternion [w, x, y, z]

        Returns:
            (roll, pitch, yaw) in radians
        """
        w, x, y, z = quat

        # Roll (x-axis rotation)
        sinr_cosp = 2 * (w * x + y * z)
        cosr_cosp = 1 - 2 * (x * x + y * y)
        roll = np.arctan2(sinr_cosp, cosr_cosp)

        # Pitch (y-axis rotation)
        sinp = 2 * (w * y - z * x)
        if abs(sinp) >= 1:
            pitch = np.copysign(np.pi / 2, sinp)  # Gimbal lock
        else:
            pitch = np.arcsin(sinp)

        # Yaw (z-axis rotation)
        siny_cosp = 2 * (w * z + x * y)
        cosy_cosp = 1 - 2 * (y * y + z * z)
        yaw = np.arctan2(siny_cosp, cosy_cosp)

        return roll, pitch, yaw

    def reset(self):
        """Reset controller state."""
        pass