import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import gradio as gr
from sklearn.neural_network import MLPRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

# -------------------------------
# Step 1: Generate synthetic data
# -------------------------------
def forward_kinematics(theta):
    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])

# Generate dataset
def create_dataset(num_samples=5000):
    X = []
    y = []
    for _ in range(num_samples):
        angles = np.random.uniform(low=-np.pi, high=np.pi, size=6)
        position = forward_kinematics(angles)
        X.append(position)
        y.append(angles)
    return np.array(X), np.array(y)

# Generate and train
X, y = create_dataset()
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

model = MLPRegressor(hidden_layer_sizes=(64, 64), max_iter=1000)
model.fit(X_train, y_train)

# Optional: Evaluate model
mse = mean_squared_error(y_test, model.predict(X_test))
print(f"Test MSE: {mse:.4f}")

# -------------------------------
# Step 2: Use ML to predict angles
# -------------------------------
def predict_angles_with_ml(target_position):
    return model.predict([target_position])[0]

def generate_trajectory(target_position):
    initial_angles = [0, 0, 0, 0, 0, 0]
    target_angles = predict_angles_with_ml(target_position)
    trajectory = np.linspace(initial_angles, target_angles, 100)
    return trajectory

def plot_trajectory(trajectory):
    x, y, z = [], [], []
    for theta in trajectory:
        pos = forward_kinematics(theta)
        x.append(pos[0])
        y.append(pos[1])
        z.append(pos[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

# -------------------------------
# Step 3: Gradio Interface
# -------------------------------
def gradio_interface(x, y, z):
    target_position = np.array([x, y, z])
    trajectory = generate_trajectory(target_position)
    fig = plot_trajectory(trajectory)
    return fig

iface = gr.Interface(
    fn=gradio_interface,
    inputs=[gr.Number(label="Target X"), gr.Number(label="Target Y"), gr.Number(label="Target Z")],
    outputs="plot",
    title="ML-Based Inverse Kinematics for 6-DOF Robot",
    description="Neural network predicts joint angles for a given (x, y, z) target using learned inverse kinematics."
)

iface.launch()