import numpy as np
import argparse

from longroad.envs import IntegerRoadEnv
import gym


import gym
from gym.spaces import Discrete, Box
from ray import tune


agentsize= 10
yellow= True
global_re1=0.01
seeds=[126291,241241516,10274]

parser = argparse.ArgumentParser(description='Trail Settings')
parser.add_argument('--agentsize',
    type=int,
    help="agentsize", default=10)
parser.add_argument('--yellow',
    type=int,
    help="yellow phase", default=1)
parser.add_argument('--seeds',
    type=int,
    help="seeds option", default=1)
parser.add_argument('--it',
    type=int,
    help="iteration option", default=500)
parser.add_argument('--path',
    help="method option", default="~/ray_results")


args=parser.parse_args()
print(args)
agentsize=args.agentsize #10
yellow=bool(args.yellow)# True
global_re1 = 0.01 if yellow else 0.1
seeds = []
iterations = args.it
if(args.seeds):
    seeds=[126291,241241516,10274]
else:
    seeds=[126291]
logpath = args.path

for seed in seeds:
    tune.run(
        "PPO",
        config={
            "env": IntegerRoadEnv,
            "num_workers": 1,
            "num_gpus": 1.0,
            "seed":seed,
            "framework": "torch",
            "env_config": {"agentsize":agentsize,"yellow":yellow,"global_re1":global_re1},
            "batch_mode": "complete_episodes",
            "observation_filter": "NoFilter",
            "logger_config": "pretty_print"},
            local_dir=logpath,
            stop={"training_iteration": 500},
            checkpoint_at_end=True)

