import numpy as np
from sim.main import InteractiveDigitalWorld
from sim.simulator import GenieSimulator, RobomimicSimulator
from sim.policy import TeleopJointPositionPolicy

if __name__ == '__main__':
    robomimic_simulator = RobomimicSimulator(env_name='lift')
    genie_simulator = GenieSimulator(
        stmaskgit_ckpt='data/genie_model/final_checkpt',
        action_stride=1,
        physics_simulator=robomimic_simulator,  # for GT image comparison
    )
    # use whatever current state is as the initial state
    genie_simulator.reset()
    teleop_policy = TeleopJointPositionPolicy(
        initial_position=robomimic_simulator.get_raw_state('robot0_joint_pos'),
        increment=1.0,
        keyboard_bindings=['1', '2', '3', '4', '5', '6', '7'],
        return_delta=True   # robomimic takes delta joint pos as action
    )
    playground = InteractiveDigitalWorld(
        simulator=genie_simulator,
        policy=teleop_policy,
        offscreen=True,
        window_size=(512 * 2, 512)  # [genie image | GT image] side-by-side
    )

    for _ in range(10):
        playground.step()

    playground.save_video(save_path='test.mp4', as_gif=False)
    playground.close()