import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import gradio as gr
from scipy.optimize import minimize

# Define forward kinematics (simplified example)
def forward_kinematics(theta):
    # Replace with actual FK for your robot
    # Example: Simple 3D position calculation
    x = np.cos(theta[0]) + np.cos(theta[1]) + np.cos(theta[2])
    y = np.sin(theta[0]) + np.sin(theta[1]) + np.sin(theta[2])
    z = theta[3] + theta[4] + theta[5]
    return np.array([x, y, z])

# Define inverse kinematics using scipy.optimize
def inverse_kinematics(target_position, initial_angles):
    def error_function(theta):
        current_position = forward_kinematics(theta)
        return np.linalg.norm(target_position - current_position)

    # Use scipy.optimize to minimize the error
    result = minimize(error_function, initial_angles, method='BFGS')
    return result.x

# Generate trajectory (from initial to target position)
def generate_trajectory(target_position):
    initial_angles = [0, 0, 0, 0, 0, 0]  # Initial joint angles
    target_angles = inverse_kinematics(target_position, initial_angles)
    # Interpolate between initial and target angles
    trajectory = np.linspace(initial_angles, target_angles, 100)
    return trajectory

# Plot the trajectory in 3D
def plot_trajectory(trajectory):
    x, y, z = [], [], []
    for theta in trajectory:
        position = forward_kinematics(theta)
        x.append(position[0])
        y.append(position[1])
        z.append(position[2])

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot(x, y, z, label="Robot Trajectory")
    ax.scatter(x[0], y[0], z[0], color='green', label="Start")
    ax.scatter(x[-1], y[-1], z[-1], color='red', label="Target")
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")
    ax.legend()
    return fig

# Gradio Interface
def gradio_interface(x, y, z):
    target_position = np.array([x, y, z])

    # Generate trajectory
    trajectory = generate_trajectory(target_position)

    # Plot trajectory
    fig = plot_trajectory(trajectory)
    return fig

# Launch Gradio App
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Number(label="Target X"),
        gr.Number(label="Target Y"),
        gr.Number(label="Target Z")
    ],
    outputs="plot",
    live=False,  # Disable live updates to show Submit button
    title="6-DOF Robot Path Planning with Inverse Kinematics",
    description="Enter the target (x, y, z) position and click Submit to generate the optimized trajectory."
)

iface.launch()