#!/usr/bin/env python

import numpy as np
from WheeledRobot import MecanumRobotEnv

# MecanumRobot kinematics
l1 = np.array([+0.5, +0.5]); f1 = np.array([+1, +1])/np.sqrt(2); d1 = np.array([+1, 0])
l2 = np.array([-0.5, +0.5]); f2 = np.array([-1, +1])/np.sqrt(2); d2 = np.array([+1, 0])
l3 = np.array([-0.5, -0.5]); f3 = np.array([-1, -1])/np.sqrt(2); d3 = np.array([-1, 0])
l4 = np.array([+0.5, -0.5]); f4 = np.array([+1, -1])/np.sqrt(2); d4 = np.array([-1, 0])
r = 0.1

# wheel equations
J11 = f1[1]/(r*d1[0]*f1[1]-r*d1[1]*f1[0])
J21 = f2[1]/(r*d2[0]*f2[1]-r*d2[1]*f2[0])
J31 = f3[1]/(r*d3[0]*f3[1]-r*d3[1]*f3[0])
J41 = f4[1]/(r*d4[0]*f4[1]-r*d4[1]*f4[0])
J12 = -f1[0]/(r*d1[0]*f1[1]-r*d1[1]*f1[0])
J22 = -f2[0]/(r*d2[0]*f2[1]-r*d2[1]*f2[0])
J32 = -f3[0]/(r*d3[0]*f3[1]-r*d3[1]*f3[0])
J42 = -f4[0]/(r*d4[0]*f4[1]-r*d4[1]*f4[0])
J13 = (-f1[1]*l1[1]-f1[0]*l1[0])/(r*d1[0]*f1[1]-r*d1[1]*f1[0])
J23 = (-f2[1]*l2[1]-f2[0]*l2[0])/(r*d2[0]*f2[1]-r*d2[1]*f2[0])
J33 = (-f3[1]*l3[1]-f3[0]*l3[0])/(r*d3[0]*f3[1]-r*d3[1]*f3[0])
J43 = (-f4[1]*l4[1]-f4[0]*l4[0])/(r*d4[0]*f4[1]-r*d4[1]*f4[0])
J = np.array([[J11, J12, J13], [J21, J22, J23], [J31, J32, J33], [J41, J42, J43]])

# forward differential kinematics
F = np.linalg.pinv(J)

def controller_mecanum(t, X_I, dX_I):
    # example controller outputs
    radius = 2; period = 10
    dX_R_des = np.array([[2*np.pi*radius/period],[0],[2*np.pi/period]])
    #dX_R_des = np.array([[1],[0],[0]])
    #dX_R_des = np.array([[0],[1],[0]])
    #dX_R_des = np.array([[0],[0],[1]])

    U = J @ dX_R_des
    return U

def run_simulation():
    """Run simulation using Gymnasium environment"""
    env = MecanumRobotEnv(render_mode="human")
    observation, _ = env.reset()
    
    print("Starting Mecanum Robot Simulation")
    print("Controller: Circular motion")
    print(f"Jacobian matrix J:\n{J}")
    
    for step in range(1000):
        # Extract controller inputs from observation
        time = observation[6]                    # Current time
        X_I = observation[:3].reshape(-1, 1)     # State [x, y, theta]
        dX_I = observation[3:6].reshape(-1, 1)   # Derivatives [dx, dy, dtheta]
        
        # Call original controller
        U = controller_mecanum(time, X_I, dX_I)
        
        # Step environment
        observation, reward, terminated, truncated, _ = env.step(U.flatten())
        
        # Render the environment
        env.render()
        
        if terminated or truncated:
            print(f"Simulation ended at step {step}")
            break
    
    input("Press Enter to close the simulation window...")
    env.close()
    print("Simulation completed")

if __name__ == "__main__":
    run_simulation()