import sys
import os
import numpy as np
from examples.development.main import ExperimentRunner
from examples.instrument import run_example_local

def main(argv=None):
    """Run ExperimentRunner locally on ray.

    To run this example on cloud (e.g. gce/ec2), use the setup scripts:
    'softlearning launch_example_{gce,ec2} examples.development <options>'.

    Run 'softlearning launch_example_{gce,ec2} --help' for further
    instructions.
    """
    # __package__ should be `development.main`
    # run_example_local(__package__, argv)
    experiment = run_example_local("examples.development", argv, diy=True)
    variant = experiment['config']
    variant['run_params']['seed'] = 100
    variant['test'] = True
    runner = experiment['run'](config=variant)
    # runner._build()
    checkpoint_dir = "/home/hesc/Data/mopo-local/logs/antcustom/antcustom_c15_m0_2_10_300_1000e3/seed:6102_0_seed=6102_2023-05-06_16-33-29du7t0rte/checkpoint_3"
    runner._restore(checkpoint_dir)
    algorithm = runner.algorithm
    evaluation_environment = algorithm._evaluation_environment
    policy = algorithm._policy
    evaluation_paths = algorithm._evaluation_paths(
        policy, evaluation_environment, stack_freq=10)
    trajectories = [evaluation_paths[i]['observations'][:,:] for i in range(len(evaluation_paths))]
    rewards = [np.sum(evaluation_paths[i]['rewards']) for i in range(len(evaluation_paths))]
    
    # trajectories = np.stack(trajectories, axis=0)
    # rewards = np.stack(rewards, axis=0)
    # data_path = "/home/hesc/Data/mopo-local/data/antcustom/2023-04-29_12-37-47/"
    # if not os.path.exists(data_path):
    #     os.makedirs(data_path)
    # trajectories.dump(data_path+"trajectories_300.np")
    # rewards.dump(data_path+"rewards_300.np")
    print("Test complete.")

if __name__ == '__main__':
    main(argv=sys.argv[1:])