from ikpy.chain import Chain
from ikpy.link import OriginLink, URDFLink
from typing import Optional
import numpy as np
import math
import time

import warnings
warnings.filterwarnings("ignore", category=UserWarning, module="ikpy")

import config as cfg
import DataTypes as dt

leg_chains = {
    0: Chain.from_urdf_file(
        cfg.urdf_path,
        base_elements=[
            "base_link",
            "leg1_coxa_joint",
            "leg1_coxa",
            "leg1_femur_joint",
            "leg1_femur",
            "leg1_tibia_joint",
            "leg1_tibia",
            "leg1_tip_joint",
            "leg1_tip",
        ],
    ),
    1: Chain.from_urdf_file(
        cfg.urdf_path,
        base_elements=[
            "base_link",
            "leg2_coxa_joint",
            "leg2_coxa",
            "leg2_femur_joint",
            "leg2_femur",
            "leg2_tibia_joint",
            "leg2_tibia",
            "leg2_tip_joint",
            "leg2_tip",
        ],
    ),
    2: Chain.from_urdf_file(
        cfg.urdf_path,
        base_elements=[
            "base_link",
            "leg3_coxa_joint",
            "leg3_coxa",
            "leg3_femur_joint",
            "leg3_femur",
            "leg3_tibia_joint",
            "leg3_tibia",
            "leg3_tip_joint",
            "leg3_tip",
        ],
    ),
    3: Chain.from_urdf_file(
        cfg.urdf_path,
        base_elements=[
            "base_link",
            "leg4_coxa_joint",
            "leg4_coxa",
            "leg4_femur_joint",
            "leg4_femur",
            "leg4_tibia_joint",
            "leg4_tibia",
            "leg4_tip_joint",
            "leg4_tip",
        ],
    ),
    4: Chain.from_urdf_file(
        cfg.urdf_path,
        base_elements=[
            "base_link",
            "leg5_coxa_joint",
            "leg5_coxa",
            "leg5_femur_joint",
            "leg5_femur",
            "leg5_tibia_joint",
            "leg5_tibia",
            "leg5_tip_joint",
            "leg5_tip",
        ],
    ),
    5: Chain.from_urdf_file(
        cfg.urdf_path,
        base_elements=[
            "base_link",
            "leg6_coxa_joint",
            "leg6_coxa",
            "leg6_femur_joint",
            "leg6_femur",
            "leg6_tibia_joint",
            "leg6_tibia",
            "leg6_tip_joint",
            "leg6_tip",
        ],
    ),
}

for leg in leg_chains.values():
    leg.active_links_mask = [False, True, True, True, False]


def ikpyForward(target_rad: dt.RadArray) -> dt.PosArray:
    target_pos_temp = []

    for leg_id, chain in leg_chains.items():
        # Forward kinematics -> 4x4 matrix
        fk_matrix = chain.forward_kinematics([0] + list(target_rad[leg_id]) + [0])
        # Extract translation vector (x, y, z)
        x, y, z = fk_matrix[:3, 3]
        target_pos_temp.append([x, y, z])
    return dt.RadArray(np.array(target_pos_temp))

def ikpyInverse(target_pos: dt.PosArray, initial_rad: Optional[dt.RadArray] = None) -> dt.RadArray:
    calculated_rads = []

    for leg_id in range(6):
        chain = leg_chains[leg_id]  # Your IKPy Chain object
        target_xyz = target_pos[leg_id]

        full_initial_position = None
        if initial_rad is not None:
            # Create a zero array matching the total number of links in the chain
            full_initial_position = [0.0] * len(chain.links)
            
            # Map the 3 active joint angles into active link indices
            active_indices = [i for i, active in enumerate(chain.active_links_mask) if active]
            
            # Match 3 active joints to the 3 active link positions in the chain
            for idx, angle in zip(active_indices, initial_rad[leg_id]):
                full_initial_position[idx] = angle

        if full_initial_position is not None:
            angles = chain.inverse_kinematics(
                target_position=target_xyz, 
                initial_position=full_initial_position
            )
        else:
            angles = chain.inverse_kinematics(target_position=target_xyz)

        # Extract only the active joint angles (3 revolute joints) from IKPy result
        active_angles = chain.active_from_full(angles)
        calculated_rads.append(active_angles)

    return dt.RadArray(calculated_rads)


def ikpytest():
    targets: dt.PosArray = dt.PosArray(
        [
            [0.047, 0.272, 0],
            [0, 0.272, 0],
            [-0.047, 0.272, 0],
            [-0.047, -0.272, 0],
            [0, -0.272, 0],
            [0.047, -0.272, 0],
        ]
    )
    joint_angles_deg: dt.DegArray = dt.DegArray(
        [
            [90, 90, 90],
            [90, 90, 90],
            [90, 90, 90],
            [90, 90, 90],
            [90, 90, 90],
            [90, 90, 90],
        ]
    )
    current_rad: dt.RadArray = joint_angles_deg.to_rad()
    ikpyInverse(targets, current_rad)