#!/usr/bin/env python

import numpy as np
from WheeledRobot import OmniRobotEnv

# OmniRobot kinematics
a1 = 0.0;         b1 = 0.0; l1 = 0.5
a2 = (2/3)*np.pi; b2 = 0.0; l2 = 0.5
a3 = (4/3)*np.pi; b3 = 0.0; l3 = 0.5
r = 0.1

# Kinematics matrix
J = np.array([
    [1/r*np.sin(a1+b1), -1/r*np.cos(a1+b1), -l1/r*np.cos(b1)],
    [1/r*np.sin(a2+b2), -1/r*np.cos(a2+b2), -l2/r*np.cos(b2)],
    [1/r*np.sin(a3+b3), -1/r*np.cos(a3+b3), -l3/r*np.cos(b3)]
])
F = np.linalg.inv(J)

def controller_omni(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 = OmniRobotEnv(render_mode="human")
    observation, _ = env.reset()
    
    print("Starting Omnidirectional Robot Simulation")
    print("Controller: Circular motion")
    
    for step in range(200):
        # 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_omni(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()