import numpy as np
import matplotlib.pyplot as plt
import gradio as gr
from scipy.spatial.distance import euclidean
from scipy.optimize import minimize

# Define the robot's arm lengths (6-DOF)
L = np.array([1, 1, 1, 1, 1, 1])  # Lengths of each arm segment

# Forward kinematics: Compute the end-effector position given joint angles
def forward_kinematics(theta):
    theta = np.array(theta)
    x = L[0] * np.cos(theta[0]) + L[1] * np.cos(theta[0] + theta[1]) + L[2] * np.cos(theta[0] + theta[1] + theta[2]) + \
         L[3] * np.cos(theta[0] + theta[1] + theta[2] + theta[3]) + L[4] * np.cos(theta[0] + theta[1] + theta[2] + theta[3] + theta[4]) + \
         L[5] * np.cos(theta[0] + theta[1] + theta[2] + theta[3] + theta[4] + theta[5])
    y = L[0] * np.sin(theta[0]) + L[1] * np.sin(theta[0] + theta[1]) + L[2] * np.sin(theta[0] + theta[1] + theta[2]) + \
         L[3] * np.sin(theta[0] + theta[1] + theta[2] + theta[3]) + L[4] * np.sin(theta[0] + theta[1] + theta[2] + theta[3] + theta[4]) + \
         L[5] * np.sin(theta[0] + theta[1] + theta[2] + theta[3] + theta[4] + theta[5])
    return np.array([x, y])

# Objective function: Minimize the distance between the end-effector and the target
def objective_function(theta, target):
    end_effector = forward_kinematics(theta)
    return euclidean(end_effector, target)

# Find the shortest path (joint angles) to reach the target
def find_shortest_path(initial_angles, target):
    result = minimize(objective_function, initial_angles, args=(target,), method='BFGS')
    return result.x

# Gradio interface
def robot_pathfinder(target_x, target_y):
    target = np.array([target_x, target_y])
    initial_angles = np.array([0, 0, 0, 0, 0, 0])  # Initial joint angles
    optimal_angles = find_shortest_path(initial_angles, target)
    
    # Plot the robot arm
    fig, ax = plt.subplots()
    ax.set_xlim(-5, 5)
    ax.set_ylim(-5, 5)
    ax.set_aspect('equal')
    
    # Draw the robot arm
    x = [0]
    y = [0]
    for i in range(len(L)):
        x.append(x[-1] + L[i] * np.cos(np.sum(optimal_angles[:i+1])))
        y.append(y[-1] + L[i] * np.sin(np.sum(optimal_angles[:i+1])))
    ax.plot(x, y, 'o-', lw=2, markersize=10)
    
    # Mark the target
    ax.plot(target[0], target[1], 'rx', markersize=10)
    
    plt.title("6-DOF Robot Arm Pathfinding")
    return fig

# Gradio app
iface = gr.Interface(
    fn=robot_pathfinder,
    inputs=["number", "number"],
    outputs="plot",
    live=True,
    title="6-DOF Robot Shortest Path Finder",
    description="Enter target coordinates (x, y) to find the shortest path for a 6-DOF robot arm."
)

iface.launch()