import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import SubprocVecEnv

import highway_env  # noqa: F401

def make_configure_env(**kwargs):
    env = gym.make(kwargs["id"], config=kwargs["config"])
    env.reset()
    return env
# ==================================
#        Main script
# ==================================

if __name__ == "__main__":
    train = True
    if train:
        n_cpu = 6
        batch_size = 64
        env_kwargs = {
        "id": "highway-v0",
        "render_mode": "human",
        "config": {
            "observation": {
            "type": "Kinematics",
            "vehicles_count": 7,
            "absolute": False,
            "order": "sorted",
            },
            "action": {
                    "type": "ContinuousAction",
            },
            "lanes_count": 3,
            "vehicles_count": 6,
            "initial_spacing": 1,
            "initial_lane_id": None,
            "controlled_vehicles": 1,
            "vehicles_density": 2,
            "simulation_frequency": 10,  # [Hz]
            "policy_frequency": 10,  # [Hz]
            "other_vehicles_type": "highway_env.vehicle.behavior.IDMVehicle",
            "centering_position": [0.8, 0.5],
            "right_lane_reward": 0,
            "policy_frequency": 2,  # [Hz]
            "show_trajectories": True,
            "offroad_terminal": True,
            # "offscreen_rendering": False
        },
    }

    env = make_configure_env(**env_kwargs)
    model = PPO(
        "MlpPolicy",
        env,
        policy_kwargs=dict(net_arch=[dict(pi=[256, 256], vf=[256, 256])]),
        n_steps=256,
        batch_size=batch_size,
        n_epochs=10,
        learning_rate=5e-4,
        gamma=0.99,
        verbose=1,
        seed = 42,
        tensorboard_log="../data/highway_ppo/",
    )
        # Train the agent
    model.learn(total_timesteps=int(2e6))
    # Save the agent
    model.save("../data/highway_ppo/model")

    model = PPO.load("../data/highway_ppo/model")
    # env = gym.make("highway-v0")
    for _ in range(5):
        obs, info = env.reset()
        done = truncated = False
        while not (done or truncated):
            action, _ = model.predict(obs)
            obs, reward, done, truncated, info = env.step(action)
            env.render()