#!/usr/bin/env python3
import sys
import os

import pinocchio as pin
import numpy as np
import osqp
from scipy import sparse
from math import radians, degrees, pi, cos, sin
import time



class KinematicsSolver():
    def __init__(self,urdf_path="urdf_rm75/RM75-B.urdf", mesh_dir="urdf_rm75"):
        """
        for realman 75b
        Initialize robotic arm kinematics using Pinocchio (ROS2 version).
        unit: m, rad
        """
        print(f' ------------  the qp based kinematic initialising   -----------')
        self.model, collision_model, visual_model = pin.buildModelsFromUrdf(urdf_path, mesh_dir)

        # -------------------------------------------------
        # ee
        # -------------------------------------------------
        ee_offset = pin.SE3(np.eye(3), np.array([0, 0, 0.0]))
        self.model.addFrame(
            pin.Frame(
                "ee",
                self.model.getJointId("joint_7"),
                self.model.getFrameId("link_7"),
                ee_offset,
                pin.FrameType.OP_FRAME
            )
        )


        # -------------------------------------------------
        # Scissor tool
        # -------------------------------------------------
        scissor_offset = pin.SE3(
            np.eye(3),
            np.array([0.0, 0.0, 0.144])
        )
        self.model.addFrame(
            pin.Frame(
                "scissor",
                self.model.getJointId("joint_7"),
                self.model.getFrameId("link_7"),
                scissor_offset,
                pin.FrameType.OP_FRAME
            )
        )

        # -------------------------------------------------
        # Camera tool
        # -------------------------------------------------
        camera_rotation = pin.rpy.rpyToMatrix(
            radians(-90),
            0,
            radians(-90)
        )
        camera_offset = pin.SE3(
            camera_rotation,
            np.array([0.05, 0.02, 0.10])
        )
        self.model.addFrame(
            pin.Frame(
                "camera",
                self.model.getJointId("joint_7"),
                self.model.getFrameId("link_7"),
                camera_offset,
                pin.FrameType.OP_FRAME
            )
        )

        # -------------------------------------------------
        # Store tool frame IDs
        # -------------------------------------------------

        self.tool_frames = {
            "scissor": self.model.getFrameId("scissor"),
            "camera": self.model.getFrameId("camera"),
            "ee":   self.model.getFrameId("ee")
        }




        self.data = self.model.createData()

        self.cfg_j_limit()

        # ----------  for reused qp_solver  ------------------
        self.nv = 7

        # Full dense symmetric matrix structure
        # P_template = np.triu(np.ones((7, 7)))
        self.P_pattern = sparse.triu(np.ones((7,7))).tocsc()

        P_sparse = sparse.csc_matrix(self.P_pattern)

        A_sparse = sparse.eye(7, format='csc')

        self.osqp_solver = osqp.OSQP()

        self.osqp_solver.setup(
            P=P_sparse,
            q=np.zeros(7),
            A=A_sparse,
            l=-np.ones(7),
            u=np.ones(7),
            verbose=False,
            warm_start=True,
            polish=False
        )

        self.W = np.diag([1, 1, 1, 0.4, 0.4, 0.4])

    def cfg_j_limit(self, min_j=None, max_j=None, rad_flag = True):
        if min_j is None:
            min_j = [-3.14159, -2.2689, -3.14159, -2.3562, -3.14159, -2.234, -6.14159]
        if max_j is None:
            max_j = [3.14159, 2.2689, 3.14159, 2.3562, 3.14159, 2.234, 6.14159]
        if rad_flag:
            for i in range(7):
                self.model.lowerPositionLimit[i] = min_j[i]
                self.model.upperPositionLimit[i] = max_j[i]
        else:
            for i in range(7):
                self.model.lowerPositionLimit[i] = min_j[i] / 180 * pi
                self.model.upperPositionLimit[i] = max_j[i] / 180 * pi

    def forward_kinematics(self, joint_angles, tool="ee"):
        """
        Compute forward kinematics.

        Args:
            joint_angles: List or array of 7 joint angles (radians)
            tool: Name of frame to compute

        Returns:
            dict: Position, rotation, rpy, quaternion
                unit:   position:   m
                        rpy:        rad
        """
        if len(joint_angles) != 7:
            raise ValueError(f"RM75 has 7 joints, got {len(joint_angles)}")

        # Create configuration vector
        q = pin.neutral(self.model)
        for i, angle in enumerate(joint_angles):
            q[i] = angle

        # Compute forward kinematics
        pin.forwardKinematics(self.model, self.data, q)
        pin.updateFramePlacements(self.model, self.data)

        # Get frame transform
        frame_id = self.tool_frames[tool]
        frame_transform = self.data.oMf[frame_id]

        # Extract results
        position = frame_transform.translation.copy()
        rotation = frame_transform.rotation.copy()

        # Compute RPY
        rpy = pin.rpy.matrixToRpy(rotation)

        # Compute quaternion
        # quat = pin.Quaternion(rotation)
        pose = np.concatenate([position, rpy], axis=0)
        return pose
        # return {
        #     'position': position,
        #     # 'rotation': rotation,
        #     'rpy': rpy,
        #     'quaternion': [quat.x, quat.y, quat.z, quat.w],
        #     # 'transform': frame_transform
        # }

    def inverse_kinematics(self, target_position, target_rpy=None,
                           target_quat=None, initial_guess=None,
                           max_iter=500, tolerance=5e-3, debug=False, tool="ee"):
        """
        Compute inverse kinematics using differential IK with multiple strategies.

        Args:
            target_position: [x, y, z] target position (meters)
            target_rpy: [roll, pitch, yaw] target orientation (radians)
            target_quat: [x, y, z, w] target orientation as quaternion
            initial_guess: Initial joint angles (radians)
            max_iter: Maximum iterations
            tolerance: Error tolerance
            debug: Print debug information
            tool: the frame name ('scissor', 'camera', 'ee')

        Returns:
            tuple: (joint_angles, success, error)
        """
        # Build target SE3 placement
        if target_quat is not None:
            quat = pin.Quaternion(target_quat[3], target_quat[0],
                                  target_quat[1], target_quat[2])
            target_rotation = quat.matrix()
        elif target_rpy is not None:
            target_rotation = pin.rpy.rpyToMatrix(target_rpy[0],
                                                  target_rpy[1],
                                                  target_rpy[2])
        else:
            target_rotation = np.eye(3)

        target_placement = pin.SE3(target_rotation, np.array(target_position))

        # Try multiple initial guesses
        initial_guesses = []

        if initial_guess is not None:
            initial_guesses.append(initial_guess)
        else:
            # Try different initial configurations
            initial_guesses.append([0.1] * 7)  # Zero config


        best_solution = None
        best_error = float('inf')

        for guess_idx, guess in enumerate(initial_guesses):
            q = pin.neutral(self.model)
            for i, angle in enumerate(guess):
                if i < len(q):
                    q[i] = np.clip(angle, self.model.lowerPositionLimit[i],
                                   self.model.upperPositionLimit[i])
            q_ref = q.copy()

            # Differential IK with adaptive damping
            damping = 0.1
            damping_reduction = 0.95
            iter_count = 0
            prev_error = float('inf')

            ee_frame_id = self.tool_frames[tool]

            J = pin.computeFrameJacobian(
                self.model,
                self.data,
                q,
                ee_frame_id,
                pin.ReferenceFrame.LOCAL
            )

            pin.forwardKinematics(self.model, self.data, q)
            pin.updateFramePlacements(self.model, self.data)

            current_placement = self.data.oMf[ee_frame_id]

            error_SE3 = current_placement.actInv(target_placement)
            error_vec = pin.log(error_SE3).vector

            # print("\n initial error =", np.linalg.norm(error_vec))
            # print(error_vec)

            while iter_count < max_iter:
                # Compute forward kinematics

                pin.computeJointJacobians(self.model, self.data, q)
                pin.framesForwardKinematics(self.model, self.data, q)

                # Get current end-effector placement
                current_placement = self.data.oMf[ee_frame_id]

                # Compute error
                error_SE3 = current_placement.actInv(target_placement)
                error_vec = pin.log(error_SE3).vector
                error_norm = np.linalg.norm(error_vec)

                if error_norm < tolerance:
                    if error_norm < best_error:
                        best_error = error_norm
                        best_solution = q[:7].copy()
                    break

                # Check if error is increasing (diverging)
                if error_norm > prev_error * 1.1 and iter_count > 10:
                    damping = min(1.0, damping * 1.5)
                else:
                    damping = max(0.01, damping * damping_reduction)


                J = pin.getFrameJacobian(
                    self.model,
                    self.data,
                    ee_frame_id,
                    pin.ReferenceFrame.LOCAL
                )

                # =========================
                # QP-based IK
                # =========================
                w_posture = 0.0001

                J_eff = pin.Jlog6(error_SE3) @ J #J #

                H = J_eff.T @ self.W @ J_eff


                # H = J.T @ self.W @ J
                H += damping * damping * np.eye(7)
                H += w_posture * np.eye(7)

                H_triu = sparse.triu(H).tocsc()

                g = -J_eff.T @ self.W @ error_vec
                g += w_posture * (q[:7] - q_ref[:7])
                # g = - J.T @ self.W @ error_vec

                # -------------------------
                # Joint velocity constraints
                # -------------------------

                dq_limit = 0.05  # rad per iteration

                lb = -dq_limit * np.ones(7)
                ub = dq_limit * np.ones(7)

                # -------------------------
                # Joint position constraints
                # -------------------------

                q_min_step = self.model.lowerPositionLimit[:7] - q[:7]
                q_max_step = self.model.upperPositionLimit[:7] - q[:7]

                lb = np.maximum(lb, q_min_step)
                ub = np.minimum(ub, q_max_step)

                # -------------------------
                # Solve QP
                # ------------------------
                # Update solver
                self.osqp_solver.update(
                    Px= H_triu.data, #H[np.triu_indices(7)], #
                    q=g,
                    l=lb,
                    u=ub
                )

                # Solve
                result = self.osqp_solver.solve()
                if result.info.status != 'solved':
                    break

                dq = result.x

                if dq is None:
                    break

                # Apply joint limits with scaling
                alpha = 1.0
                q = pin.integrate(self.model, q, alpha * dq)

                prev_error = error_norm
                iter_count += 1

        if best_solution is not None:
            # return best_solution, True, best_error, iter_count
            return 0, best_solution
        else:
            # return q[:7].copy(), False, error_norm, iter_count
            return -1, q[:7].copy()

    # def invese_kinematics_velocity(self, target_position, target_rpy=None,
    #                        target_quat=None, initial_guess=None, tool="ee"):
    #     """
    #             Compute the converging velocity (motion direction) of joints based on qp inverse kinematics.
    #
    #             Args:
    #                 target_position: [x, y, z] target position (meters)
    #                 target_rpy: [roll, pitch, yaw] target orientation (radians)
    #                 target_quat: [x, y, z, w] target orientation as quaternion
    #                 initial_guess: Initial joint angles (radians)
    #                 tool: the frame name ('scissor', 'camera', 'ee')
    #
    #             Returns:
    #                 joint_velocity: np.array()
    #             """
    #     # Build target SE3 placement
    #     if target_quat is not None:
    #         quat = pin.Quaternion(target_quat[3], target_quat[0],
    #                               target_quat[1], target_quat[2])
    #         target_rotation = quat.matrix()
    #     elif target_rpy is not None:
    #         target_rotation = pin.rpy.rpyToMatrix(target_rpy[0],
    #                                               target_rpy[1],
    #                                               target_rpy[2])
    #     else:
    #         target_rotation = np.eye(3)
    #
    #     target_placement = pin.SE3(target_rotation, np.array(target_position))
    #
    #     # Try multiple initial guesses
    #     initial_guesses = []
    #
    #     if initial_guess is not None:
    #         initial_guesses.append(initial_guess)
    #     else:
    #         # Try different initial configurations
    #         initial_guesses.append([0.1] * 7)  # Zero config
    #         initial_guesses.append([radians(30), radians(45), radians(30),
    #                                 radians(-45), radians(30), radians(-30), 0])
    #         initial_guesses.append([radians(-30), radians(45), radians(-30),
    #                                 radians(45), radians(30), radians(30), 0])
    #
    #     best_solution = None
    #     best_error = float('inf')
    #
    #     for guess_idx, guess in enumerate(initial_guesses):
    #         q = pin.neutral(self.model)
    #         for i, angle in enumerate(guess):
    #             if i < len(q):
    #                 q[i] = np.clip(angle, self.model.lowerPositionLimit[i],
    #                                self.model.upperPositionLimit[i])
    #
    #         # Differential IK with adaptive damping
    #         damping = 0.01
    #         damping_reduction = 0.95
    #         iter_count = 0
    #         prev_error = float('inf')
    #
    #         ee_frame_id = self.tool_frames[tool]
    #
    #         J = pin.computeFrameJacobian(
    #             self.model,
    #             self.data,
    #             q,
    #             ee_frame_id,
    #             pin.ReferenceFrame.LOCAL_WORLD_ALIGNED
    #         )
    #
    #         while iter_count < max_iter:
    #             # Compute forward kinematics
    #
    #             pin.computeJointJacobians(self.model, self.data, q)
    #             pin.framesForwardKinematics(self.model, self.data, q)
    #
    #             # Get current end-effector placement
    #
    #             current_placement = self.data.oMf[ee_frame_id]
    #
    #             # Compute error
    #             error_SE3 = current_placement.actInv(target_placement)
    #             error_vec = pin.log(error_SE3).vector
    #             error_norm = np.linalg.norm(error_vec)
    #
    #             if error_norm < tolerance:
    #                 joint_angles = q[:7].copy()
    #                 fk_result = self.forward_kinematics(joint_angles, tool=tool)
    #                 position_error = np.linalg.norm(fk_result['position'] - np.array(target_position))
    #
    #                 if position_error < best_error:
    #                     best_error = position_error
    #                     best_solution = joint_angles
    #                 break
    #
    #             # Check if error is increasing (diverging)
    #             if error_norm > prev_error * 1.1 and iter_count > 10:
    #                 damping = min(1.0, damping * 1.5)
    #             else:
    #                 damping = max(0.01, damping * damping_reduction)
    #
    #             J = pin.getFrameJacobian(
    #                 self.model,
    #                 self.data,
    #                 ee_frame_id,
    #                 pin.ReferenceFrame.LOCAL_WORLD_ALIGNED
    #             )
    #
    #             # =========================
    #             # QP-based IK
    #             # =========================
    #
    #             H = J.T @ self.W @ J
    #             H += damping * damping * np.eye(7)
    #
    #             H_triu = sparse.triu(H).tocsc()
    #
    #             g = -J.T @ self.W @ error_vec
    #
    #             # -------------------------
    #             # Joint velocity constraints
    #             # -------------------------
    #
    #             dq_limit = 0.05  # rad per iteration
    #
    #             lb = -dq_limit * np.ones(7)
    #             ub = dq_limit * np.ones(7)
    #
    #             # -------------------------
    #             # Joint position constraints
    #             # -------------------------
    #
    #             q_min_step = self.model.lowerPositionLimit[:7] - q[:7]
    #             q_max_step = self.model.upperPositionLimit[:7] - q[:7]
    #
    #             lb = np.maximum(lb, q_min_step)
    #             ub = np.minimum(ub, q_max_step)
    #
    #             # -------------------------
    #             # Solve QP
    #             # ------------------------
    #             # Update solver
    #             self.osqp_solver.update(
    #                 Px=H_triu.data,
    #                 q=g,
    #                 l=lb,
    #                 u=ub
    #             )
    #
    #             # Solve
    #             result = self.osqp_solver.solve()
    #
    #             if result.info.status != 'solved':
    #                 break
    #
    #             dq = result.x
    #
    #             if dq is None:
    #                 break
    #
    #             # Apply joint limits with scaling
    #             alpha = 0.5
    #             q = pin.integrate(self.model, q, alpha * dq)
    #
    #             prev_error = error_norm
    #             iter_count += 1
    #
    #     if best_solution is not None:
    #         return best_solution, True, best_error
    #     else:
    #         return None, False, None

    def compute_jacobian(self, joint_angles, tool="ee"):
        """Compute geometric Jacobian (6x7)"""
        q = pin.neutral(self.model)
        for i, angle in enumerate(joint_angles):
            q[i] = angle

        pin.forwardKinematics(self.model, self.data, q)
        pin.updateFramePlacements(self.model, self.data)
        ee_frame_id = self.tool_frames[tool]
        J = pin.computeFrameJacobian(self.model, self.data, q, ee_frame_id)

        return J

    def get_subchain_jacobian(self,
                              joint_angles,
                              frame_names
                              ):

        q = pin.neutral(self.model)

        all_active_joints = self.get_active_joints_from_frame(frame_names)

        for i in range(7):
            q[i] = joint_angles[i]

        pin.forwardKinematics(self.model, self.data, q)
        pin.updateFramePlacements(self.model, self.data)
        pin.computeJointJacobians(self.model, self.data, q)

        Js = []

        for  frame_name, active_joints in zip(frame_names, all_active_joints):
            frame_id = self.model.getFrameId(frame_name)

            J = pin.getFrameJacobian(
                self.model,
                self.data,
                frame_id,
                pin.ReferenceFrame.LOCAL
            )
            Js.append(J[:, active_joints])

        return Js

    def get_active_joints_from_frame(self, frame_names):
        """
        Return active joint indices affecting a frame.

        Example:
            frame_name='link_4'
            -> [0,1,2,3]
        """
        all_active_joint_ids = []
        for frame_name in frame_names:
            frame_id = self.model.getFrameId(frame_name)

            # Parent joint of this frame
            joint_id = self.model.frames[frame_id].parentJoint

            print(f'frame_id = {frame_id}, and joint_id = {joint_id}')

            active_joint_ids = []


            # Traverse upward to root
            while joint_id > 0:
                # Pinocchio joint indexing:
                # universe joint = 0
                # robot joints start from 1

                active_joint_ids.append(joint_id - 1)

                # Move to parent joint
                joint_id = self.model.parents[joint_id]

            # Reverse so order becomes base -> tip
            active_joint_ids.reverse()
            all_active_joint_ids.append(active_joint_ids)

        return all_active_joint_ids

    def plan_cartesian_trajectory(self, start_pos, end_pos,
                                  start_rpy=None, end_rpy=None,
                                  num_steps=20, tool='ee'):
        """
        Plan a Cartesian trajectory with IK for each waypoint.
        """
        # Get current end-effector pose if start_rpy not provided
        if start_rpy is None:
            # Try to find a valid starting configuration
            test_angles = [0.1] * 7
            fk_test = self.forward_kinematics(test_angles,tool=tool)
            start_rpy = fk_test['rpy']

        if end_rpy is None:
            end_rpy = start_rpy

        # First, check if target is reachable
        print(f"\nChecking if target is reachable...")
        target_pos = end_pos
        target_rpy = end_rpy

        test_solution, success, error = self.inverse_kinematics(
            target_pos, target_rpy=target_rpy, initial_guess=[0.1] * 7, max_iter=500, tool=tool
        )

        if not success:
            print(f"Warning: Target may be unreachable or difficult to reach")
            print(f"Trying with relaxed tolerance...")

        # Initial guess for IK (start with zero configuration)
        current_angles = [0.1] * 7
        trajectory = []

        print(f"\nPlanning trajectory from ({start_pos[0]:.2f}, {start_pos[1]:.2f}, {start_pos[2]:.2f})")
        print(f"To ({end_pos[0]:.2f}, {end_pos[1]:.2f}, {end_pos[2]:.2f})")
        print("-" * 60)

        for i in range(num_steps + 1):
            t = i / num_steps

            # Interpolate position
            pos = [
                start_pos[0] + t * (end_pos[0] - start_pos[0]),
                start_pos[1] + t * (end_pos[1] - start_pos[1]),
                start_pos[2] + t * (end_pos[2] - start_pos[2])
            ]

            # Interpolate orientation
            rpy = [
                start_rpy[0] + t * (end_rpy[0] - start_rpy[0]),
                start_rpy[1] + t * (end_rpy[1] - start_rpy[1]),
                start_rpy[2] + t * (end_rpy[2] - start_rpy[2])
            ]

            # Compute IK
            joint_angles, success, error = self.inverse_kinematics(
                pos, target_rpy=rpy, initial_guess=current_angles, max_iter=300, tool=tool
            )

            if not success:
                print(f"  Waypoint {i}: IK failed!")
                break

            # Verify
            fk_verify = self.forward_kinematics(joint_angles, tool=tool)

            trajectory.append({
                'step': i,
                't': t,
                'position': pos,
                'rpy': rpy,
                'joint_angles': joint_angles,
                'actual_position': fk_verify['position'],
                'error': error
            })

            # Update current angles for next iteration
            current_angles = joint_angles

            if i % 5 == 0 or i == num_steps:
                print(f"  Waypoint {i:3d}: pos=({pos[0]:.3f}, {pos[1]:.3f}, {pos[2]:.3f}), "
                      f"error={error:.6f}m")

        return trajectory



def main():
    """Main test function"""


    rm75 = KinematicsSolver()

    # Test 1: Forward Kinematics
    print("\n1. Forward Kinematics Test")
    print("-" * 40)

    tool_name = "scissor"
    joint_angles_zero = [0.1] * 7
    fk_result = rm75.forward_kinematics(joint_angles_zero, tool=tool_name)

    print(f"Init configuration:")
    print(f"  Position: ({fk_result['position'][0]:.3f}, "
          f"{fk_result['position'][1]:.3f}, {fk_result['position'][2]:.3f}) m")

    # Test 2: Inverse Kinematics with more reachable target
    print("\n2. Inverse Kinematics Test")
    print("-" * 40)

    # Try a simpler target first
    target_pos = [0.3, 0.2, 0.4]  # More reachable position
    target_rpy = [0.0, 0.0, radians(45)]  # Simpler orientation

    print(f"Target: ({target_pos[0]:.3f}, {target_pos[1]:.3f}, {target_pos[2]:.3f}) m")

    import time
    init_joints = [0.2] * 7
    time0 = time.time()
    for ii in range(100):
        joint_solution, success, error = rm75.inverse_kinematics(
            target_pos, target_rpy=target_rpy, initial_guess=init_joints,
            max_iter=500, debug=False, tool=tool_name
        )
    time1 = time.time()
    print(f"Time: {time1 - time0}")

    if success:
        print(f"✓ Solution found! Error: {error:.6f} m")
        for i, angle in enumerate(joint_solution):
            print(f"  Joint {i + 1}: {degrees(angle):7.2f}°")

        # Verify
        fk_verify = rm75.forward_kinematics(joint_solution,tool=tool_name)
        print(
            f"  Position: ({fk_verify['position'][0]:.3f}, {fk_verify['position'][1]:.3f}, {fk_verify['position'][2]:.3f}) m")
    else:
        print("✗ IK failed to find a solution!")

    # Test 3: Jacobian
    print("\n3. Jacobian Matrix")
    print("-" * 40)

    J = rm75.compute_jacobian(joint_angles_zero, tool=tool_name)
    print(f"Jacobian shape: {J.shape}")
    for i in range(min(3, J.shape[0])):
        row_str = "  ".join([f"{J[i, j]:7.3f}" for j in range(7)])
        print(f"  Row {i + 1}: {row_str}")

    # Test 4: Trajectory Planning with reachable positions
    print("\n4. Cartesian Trajectory Planning")
    print("-" * 40)

    start_pos = [0.3, 0.0, 0.4]  # Start position
    end_pos = [0.3, 0.0, 0.55]  # End position (smaller movement)

    fk0 = rm75.forward_kinematics([0.1] * 7, tool=tool_name)

    trajectory = rm75.plan_cartesian_trajectory(
        start_pos,
        end_pos,
        start_rpy=fk0['rpy'],
        end_rpy=[
            fk0['rpy'][0] + radians(10),
            fk0['rpy'][1],
            fk0['rpy'][2]
        ],
        num_steps=10,
        tool=tool_name
    )

    if trajectory:
        print(f"\n✓ Generated {len(trajectory)} waypoints")

    if success:
        print("✓ Inverse kinematics working (with simplified target)")
    else:
        print("⚠ Inverse kinematics may need tuning - try different targets")


    print("\n" + "=" * 60)
    print(f'test subchain Jacobian, for future obstacle avoidance')
    frame_names = [
        "link_2",
        "link_4",
        "link_7"
    ]
    Js_sub = rm75.get_subchain_jacobian(
        joint_angles=joint_angles_zero,
        frame_names=frame_names
    )
    print(f'Js_sub: {Js_sub}')

    return rm75, trajectory


if __name__ == "__main__":
    rm75, trajectory = main()

    print("\n" + "=" * 60)
    print("All tests completed!")
    print("=" * 60)