import time
from sim.policy import DiffusionPolicy
from sim.robomimic.robomimic_runner import RolloutRunner


NUM_EVAL_TRIALS = 10

if __name__ == '__main__':
    
    success_rates = [0.38, 0.52, 0.70, 1.00]
    eval_time_taken = [0.0] * len(success_rates)

    env_runner = RolloutRunner(
        env_names=['lift'],
        episode_num=NUM_EVAL_TRIALS,
        save_video=False
    )
    for index, sr in enumerate(success_rates): 
        print(f"Running evaluation for model with success rate: {sr:.2f}")
        diffusion_policy = DiffusionPolicy(f'data/dp_ckpt/dp_lift_sr{sr:.2f}.ckpt')
        n_obs_steps = diffusion_policy.n_obs_steps

        start_time = time.time()
        success, reward = env_runner.run(
            policy=diffusion_policy,
            env_name=['lift'],
        )
        end_time = time.time()
        eval_time_taken[index] = (end_time - start_time) / NUM_EVAL_TRIALS
        print(f"Time taken for evaluation: {eval_time_taken[index]}")

        print(f"success: {success}, reward: {reward}")