import json

from pgdrive.envs.pgdrive_env import PGDriveEnv
from pgdrive.scene_creator.map import Map, MapGenerateMethod
from pgdrive.scene_manager.traffic_manager import TrafficMode
from pgdrive.utils import setup_logger


class TestEnv(PGDriveEnv):
    def __init__(self, save_episode=True, vis=True):
        super(TestEnv, self).__init__(
            {
                "environment_num": 1,
                "traffic_density": 0.1,
                "start_seed": 5,
                "manual_control": vis,
                "use_render": vis,
                "traffic_mode": TrafficMode.Hybrid,
                "record_episode": save_episode,
                "map_config": {
                    Map.GENERATE_METHOD: MapGenerateMethod.BIG_BLOCK_SEQUENCE,
                    Map.GENERATE_PARA: "XTXTXTXTXT",
                    Map.LANE_WIDTH: 3.5,
                    Map.LANE_NUM: 3,
                }
            }
        )


def test_save_episode(vis=False):
    setup_logger(True)

    test_dump = False

    env = TestEnv(vis=vis)
    try:
        o = env.reset()
        epi_info = None
        for i in range(1, 100000 if vis else 2000):
            o, r, d, info = env.step([0, 1])
            if vis:
                env.render()
            if d:
                epi_info = env.scene_manager.dump_episode()

                # test dump json
                if test_dump:
                    with open("test.json", "w") as f:
                        json.dump(epi_info, f)
                break

        o = env.reset(epi_info)
        for i in range(1, 100000 if vis else 2000):
            o, r, d, info = env.step([0, 1])
            if vis:
                env.render()
            if d:
                break
    finally:
        env.close()


if __name__ == "__main__":
    test_save_episode(vis=True)
