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

if __name__ == '__main__':
    robomimic_simulator = RobomimicSimulator(env_name='lift')
    genie_simulator = GenieSimulator(
        image_encoder_type="magvit",
        image_encoder_ckpt="data/magvit2.ckpt",
        quantize=True,
        backbone_type="stmaskgit",
        backbone_ckpt="data/genie_lang/maskgit256",
        prompt_horizon=8,

        # image_encoder_type='temporalvae',
        # image_encoder_ckpt='stabilityai/stable-video-diffusion-img2vid',
        # quantize=False,
        # backbone_type="stmar",
        # backbone_ckpt="data/genie_lang/mar1024",
        # prompt_horizon=11,
        
        action_stride=1,
        domain='robomimic',
        physics_simulator=robomimic_simulator,
        compute_psnr=False,
        compute_delta_psnr=False,
    )
    # use whatever current state is as the initial state
    genie_simulator.reset()
    random_policy = RandomJointPositionPolicy(
        action_bounds=(
            np.array([-0.5] * 7),
            np.array([0.5] * 7)
        )
    )
    playground = InteractiveDigitalWorld(
        simulator=genie_simulator,
        policy=random_policy,
        offscreen=True,
        window_size=(512 * 2, 512)  # [genie image | GT image] side-by-side
    )

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

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