import os.path
from ray.tune import grid_search
from metadrive.envs.scenario_env import ScenarioEnv
from metadrive.envs.gym_wrapper import GymEnvWrapper
from scenarionet import SCENARIONET_REPO_PATH, SCENARIONET_DATASET_PATH
from scenarionet_training.train_utils.multi_worker_PPO import MultiWorkerPPO
from scenarionet_training.train_utils.utils import train, get_train_parser, get_exp_name

config = dict(
    env=GymEnvWrapper,
    env_config=dict(inner_class=ScenarioEnv, inner_config=dict(
        # scenario
        start_scenario_index=0,
        num_scenarios=40000,
        data_directory=os.path.join(SCENARIONET_DATASET_PATH, "pg_train"),
        sequential_seed=True,

        # curriculum training
        curriculum_level=100,
        target_success_rate=0.8,
        # episodes_to_evaluate_curriculum=400,  # default=num_scenarios/curriculum_level

        # traffic & light
        reactive_traffic=False,
        no_static_vehicles=True,
        no_light=True,
        static_traffic_object=True,

        # training scheme
        horizon=None,
        steering_range_penalty=2,
        heading_penalty=1.0,
        lateral_penalty=1.0,
        no_negative_reward=True,
        on_lane_line_penalty=0,
        crash_vehicle_penalty=2,
        crash_human_penalty=2,
        out_of_road_penalty=2,
        max_lateral_dist=2,
        # crash_vehicle_done=True,

        vehicle_config=dict(side_detector=dict(num_lasers=0))

    )),

    # ===== Evaluation =====
    evaluation_interval=15,
    evaluation_num_episodes=1000,
    # TODO (LQY), this is a sample from testset do eval on all scenarios after training!
    evaluation_config=dict(env_config=dict(start_scenario_index=0,
                                           num_scenarios=1000,
                                           sequential_seed=True,
                                           curriculum_level=1,  # turn off
                                           data_directory=os.path.join(SCENARIONET_DATASET_PATH, "pg_test"))),
    evaluation_num_workers=10,
    metrics_smoothing_episodes=10,

    # ===== Training =====
    model=dict(fcnet_hiddens=[512, 256, 128]),
    horizon=600,
    num_sgd_iter=20,
    lr=1e-4,
    rollout_fragment_length=500,
    sgd_minibatch_size=200,
    train_batch_size=50000,
    num_gpus=0.5,
    num_cpus_per_worker=0.3,
    num_cpus_for_driver=1,
    num_workers=20,
    framework="tf"
)

if __name__ == '__main__':
    # PG data is generated with seeds 10,000 to 60,000
    args = get_train_parser().parse_args()
    exp_name = get_exp_name(args)
    stop = int(100_000_000)
    config["num_gpus"] = 0.5 if args.num_gpus != 0 else 0

    train(
        MultiWorkerPPO,
        exp_name=exp_name,
        save_dir=os.path.join(SCENARIONET_REPO_PATH, "experiment"),
        keep_checkpoints_num=5,
        stop=stop,
        config=config,
        num_gpus=args.num_gpus,
        # num_seeds=args.num_seeds,
        num_seeds=5,
        test_mode=args.test,
        # local_mode=True,
        # TODO remove this when we release our code
        # wandb_key_file="~/wandb_api_key_file.txt",
        wandb_project="scenarionet",
    )