import numpy as np
import plotly.graph_objects as go
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])
    z = 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, z])

# 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_z):
    target = np.array([target_x, target_y, target_z])
    initial_angles = np.array([0, 0, 0, 0, 0, 0])  # Initial joint angles
    optimal_angles = find_shortest_path(initial_angles, target)
    
    # Compute joint positions for visualization
    joint_positions = []
    joint_positions.append([0, 0, 0])  # Base position
    for i in range(len(L)):
        x = joint_positions[-1][0] + L[i] * np.cos(np.sum(optimal_angles[:i+1]))
        y = joint_positions[-1][1] + L[i] * np.sin(np.sum(optimal_angles[:i+1]))
        z = joint_positions[-1][2] + L[i] * np.sin(np.sum(optimal_angles[:i+1]))
        joint_positions.append([x, y, z])
    
    # Create 3D plot using Plotly
    fig = go.Figure()
    
    # Add robot arm segments
    for i in range(len(joint_positions) - 1):
        fig.add_trace(go.Scatter3d(
            x=[joint_positions[i][0], joint_positions[i+1][0]],
            y=[joint_positions[i][1], joint_positions[i+1][1]],
            z=[joint_positions[i][2], joint_positions[i+1][2]],
            mode='lines+markers',
            line=dict(color='blue', width=5),
            marker=dict(size=5)
        ))
    
    # Add target point
    fig.add_trace(go.Scatter3d(
        x=[target[0]],
        y=[target[1]],
        z=[target[2]],
        mode='markers',
        marker=dict(color='red', size=10)
    ))
    
    # Set layout
    fig.update_layout(
        scene=dict(
            xaxis=dict(range=[-5, 5]),
            yaxis=dict(range=[-5, 5]),
            zaxis=dict(range=[-5, 5]),
            aspectmode='cube'
        ),
        title="6-DOF Robot Arm Pathfinding in 3D"
    )
    
    return fig

# Gradio app
iface = gr.Interface(
    fn=robot_pathfinder,
    inputs=[
        gr.Number(label="Target X-coordinate"),
        gr.Number(label="Target Y-coordinate"),
        gr.Number(label="Target Z-coordinate")
    ],
    outputs=gr.Plot(),  # Use Plotly for 3D visualization
    live=True,
    title="6-DOF Robot Shortest Path Finder in 3D",
    description="Enter target coordinates (x, y, z) to find the shortest path for a 6-DOF robot arm in 3D space."
)

# Launch the app
if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)