import os
import numpy as np
from tray_sim import SIMULATOR_MODEL, run_tray_simulation, classify_stability

def main():
    sim = SIMULATOR_MODEL()
    sim.create_pusher('3.0 3.0 0.05')
    
    # Tray parameters
    physical_parameters_for_object_id_tray = {
        'sliding_friction': 0.1,
        'armature': 0.1,
        'stiffness': 0.0,
        'mass': 0.5,
        'damping': 20
    }
    sim.create_tray(object_physics=physical_parameters_for_object_id_tray)
    
    # Object parameters
    physical_parameters_for_object_id_1 = {  # Bottle
        'sliding_friction': 0.22,
        'armature': 0.41,
        'stiffness': 0.53,
        'mass': 20.0,
        'damping': 7.0
    }
    physical_parameters_for_object_id_2 = {  # Martini glass
        'sliding_friction': 0.52,
        'armature': 0.63,
        'stiffness': 1.22,
        'mass': 10.0,
        'damping': 8.0
    }
    physical_parameters_for_object_id_3 = {  # Wine glass
        'sliding_friction': 0.73,
        'armature': 1.04,
        'stiffness': 0.94,
        'mass': 4.0,
        'damping': 7.9
    }
    
    # Create objects
    sim.create_object(
        object_id=1,
        object_name='bottle',
        object_location=('row_1', 'column_3'),
        object_color='orange',
        object_physics=physical_parameters_for_object_id_1
    )
    sim.create_object(
        object_id=2,
        object_name='martini_glass',
        object_location=('row_1', 'column_2'),
        object_color='orange',
        object_physics=physical_parameters_for_object_id_2
    )
    sim.create_object(
        object_id=3,
        object_name='wine_glass',
        object_location=('row_1', 'column_1'),
        object_color='orange',
        object_physics=physical_parameters_for_object_id_3
    )
    
    # Run simulation and collect trajectories
    model = sim.model
    data = sim.data
    trajectories = {
        'bottle': [],
        'martini_glass': [],
        'wine_glass': []
    }
    sample_steps = [0, 25, 50, 75, 100, 125, 150, 175, 199]  # 9 points as in PAGE13
    
    for t in range(200):
        if t == 50:  # PUSH_START_STEP
            pusher_idx = 3 * 7  # 3 objects * 7 DOF
            pusher_vel_idx = 3 * 6  # 3 objects * 6 velocities
            direction = np.array([0.0, 0.0, 0.05]) - np.array([3.0, 3.0, 0.05])
            direction[2] = 0.0
            direction /= np.linalg.norm(direction)
            push_velocity = -4.8 * direction / np.sqrt(2)  # Normalize to match (-4.8, -4.8)
            data.qvel[pusher_vel_idx:pusher_vel_idx+3] = push_velocity
        
        mujoco.mj_step(model, data)
        
        if t in sample_steps:
            for i, obj_name in enumerate(['bottle', 'martini_glass', 'wine_glass'], 1):
                pos = data.qpos[(i-1)*7:(i-1)*7+3]
                # Adjust z to match document's center of gravity
                if obj_name == 'bottle':
                    pos[2] += 1.1
                else:
                    pos[2] += 0.5
                trajectories[obj_name].append([round(p, 1) for p in pos])
    
    sim.create_scene()
    sim_out = sim.run_simulation()  # Ensure scene is finalized
    del sim
    
    # Format and save trajectories
    with open('object_traj_example_1.txt', 'w') as f:
        f.write("bottle_motion_trajectory (x, y, z) = [\n")
        for pos in trajectories['bottle']:
            f.write(f"({pos[0]}, {pos[1]}, {pos[2]}),\n")
        f.write("]\n")
        f.write("martini_glass_motion_trajectory (x, y, z) = [\n")
        for pos in trajectories['martini_glass']:
            f.write(f"({pos[0]}, {pos[1]}, {pos[2]}),\n")
        f.write("]\n")
        f.write("wine_glass_motion_trajectory (x, y, z) = [\n")
        for pos in trajectories['wine_glass']:
            f.write(f"({pos[0]}, {pos[1]}, {pos[2]}),\n")
        f.write("]\n")
    
    return trajectories

if __name__ == "__main__":
    trajectories = main()
    print("Trajectories saved to object_traj_example_1.txt")