from ikpy.chain import Chain
from ikpy.link import OriginLink, URDFLink
import matplotlib.pyplot as plt
import numpy as np
import math
import time

import GlobalVariables as gv
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) -> dt.RadArray:
    target_rad_temp = []
    print("target_pos:\n",target_pos)
    for leg_id, chain in leg_chains.items():
        # prepare initial guess
        guess = np.array([0] + list(gv.current_rad[leg_id]) + [0], dtype=float)

        try:
            ik_result = chain.inverse_kinematics(
                target_pos[leg_id],
                initial_position=guess,
                max_iter=100
            )
        except ValueError:
            # fallback → keep the clipped guess
            ik_result = guess

        # keep only the 3 actuated joint angles (indices 1,2,3)
        target_rad_temp.append(ik_result[1:4])

    return dt.RadArray(np.array(target_rad_temp))
    # for leg_id, chain in leg_chains.items():
    #     ik_result = chain.inverse_kinematics(target_pos[leg_id])
    #     # Remove first element (IKPy adds a "dummy" fixed base joint)
    #     target_rad_temp.append(ik_result[1:4])
    return dt.RadArray(np.array(target_rad_temp))


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)


# walk old
"""
def walk(duration=standard_duration, ticks=standard_tickrate, curve_height=gv.step_height):
    tick_duration = duration / ticks
    tick_positions = np.copy(gv.current_pos)
    dirmov_copy = gv.vector_dirmov 
    target_pos =[]
    for i in range(6): # Zielposition berechnen 
        target_pos.append([gv.center_points[i][0] + dirmov_copy[0] * gv.step_length, gv.center_points[i][1] + dirmov_copy[1] * gv.step_length, gv.robot_height ])
    def interpolate(t, p0, p1, p2):
        return (1 - t)**2 * p0 + 2 * (1 - t) * t * p1 + t**2 * p2

    for tick in range(int(ticks) + 1):
        loop_start = time.perf_counter()
        t = tick / ticks

        for leg_id in range(6):
            if gv.leg_state[leg_id] == "drag":
                # Linear interpolation
                tick_positions[leg_id] = gv.current_pos[leg_id] + (gv.center_points[leg_id] - gv.current_pos[leg_id]) * t

            elif gv.leg_state[leg_id] == "step":
                # Curve movement (Bezier path)
                help_pos = [
                    target_pos[leg_id][0] - gv.current_pos[leg_id][0],
                    target_pos[leg_id][1] - gv.current_pos[leg_id][1],
                    target_pos[leg_id][2] + curve_height
                ]
                x = interpolate(t, gv.current_pos[leg_id][0], help_pos[0], target_pos[leg_id][0])
                y = interpolate(t, gv.current_pos[leg_id][1], help_pos[1], target_pos[leg_id][1])
                z = interpolate(t, gv.current_pos[leg_id][2], help_pos[2], target_pos[leg_id][2])
                tick_positions[leg_id] = [x, y, z]

        # Send all legs at once through IK
        ikpyInverse(tick_positions, gv.current_pos)

        elapsed = time.perf_counter() - loop_start
        sleep_time = tick_duration - elapsed
        if sleep_time > 0:
            time.sleep(sleep_time)

    # Final position correction
    ikpyInverse(target_pos, gv.current_pos)

    if (gv.leg_state[0] == "step"):
        gv.leg_state = np.array(["drag", "step", "drag", "step", "drag", "step"])
    if (gv.leg_state[0] == "drag"):
        gv.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
    time.sleep(0.05)



# Berechnung der Bein Bewegung - Grade
def drag_leg(
    leg_num,
    current_pos,
    target_pos,
    duration=cfg.standard_duration,
    ticks=cfg.standard_tickrate,
):
    tick_duration = duration / ticks
    dragtickvec = (target_pos - current_pos[leg_num]) / ticks

    tick_positions = np.copy(current_pos)

    for tick in range(int(ticks)):
        start_time = time.perf_counter()

        tick_positions[leg_num] += dragtickvec
        ikpyInverse(tick_positions, current_pos)

        elapsed = time.perf_counter() - start_time
        sleep_time = tick_duration - elapsed
        if sleep_time > 0:
            time.sleep(sleep_time)

    tick_positions[leg_num] = target_pos
    ikpyInverse(tick_positions, current_pos)
    time.sleep(0.05)


# Berechnung der Bein Bewegung - Kurve
def curve_leg(
    leg_num,
    current_pos,
    target_pos,
    curve_height=cfg.step_height,
    duration=cfg.standard_duration,
    ticks=cfg.standard_tickrate,
):
    tick_duration = duration / ticks

    help_pos = [
        target_pos[0] - current_pos[leg_num][0],
        target_pos[1] - current_pos[leg_num][1],
        target_pos[2] + curve_height,
    ]

    def interpolate(t, p0, p1, p2):
        return (1 - t) ** 2 * p0 + 2 * (1 - t) * t * p1 + t**2 * p2

    tick_positions = np.copy(current_pos)

    for tick in range(int(ticks) + 1):
        start_time = time.perf_counter()

        t = tick / ticks
        x = interpolate(t, current_pos[leg_num][0], help_pos[0], target_pos[0])
        y = interpolate(t, current_pos[leg_num][1], help_pos[1], target_pos[1])
        z = interpolate(t, current_pos[leg_num][2], help_pos[2], target_pos[2])

        tick_positions[leg_num] = [x, y, z]
        ikpyInverse(tick_positions, current_pos)

        elapsed = time.perf_counter() - start_time
        sleep_time = tick_duration - elapsed
        if sleep_time > 0:
            time.sleep(sleep_time)

    tick_positions[leg_num] = target_pos
    ikpyInverse(tick_positions, current_pos)
    time.sleep(0.05)

"""
if __name__ == "__main__":
    print(ikpyForward(gv.test_deg.to_rad()))