from navsim.visualization.plots import plot_bev_frame, plot_cameras_frame

from hydra.core.global_hydra import GlobalHydra
import os
from pathlib import Path

import hydra
from hydra.utils import instantiate
import numpy as np
import matplotlib.pyplot as plt

from navsim.common.dataloader import SceneLoader
from navsim.common.dataclasses import SceneFilter, SensorConfig
GlobalHydra.instance().clear()
# os.environ['OPENSCENE_DATA_ROOT'] = '/mnt/g/navsim/'
# os.environ['NUPLAN_MAPS_ROOT'] = '/mnt/g/navsim/maps'
# os.environ['NUPLAN_MAP_VERSION'] = "nuplan-maps-v1.0"
# os.environ['NAVSIM_EXP_ROOT'] = '/mnt/g/navsim_exp'
# os.environ['NAVSIM_DEVKIT_ROOT'] = '/mnt/f/e2e/navsim_ours'

SPLIT = "tiny"  # ["mini", "test", "trainval"]
FILTER = "all_scenes"

hydra.initialize(config_path="../planning/script/config/common/scene_filter")
cfg = hydra.compose(config_name=FILTER)
scene_filter: SceneFilter = instantiate(cfg)
openscene_data_root = Path(os.getenv("OPENSCENE_DATA_ROOT"))

scene_loader = SceneLoader(
    openscene_data_root / f"navsim_logs/{SPLIT}",
    openscene_data_root / f"sensor_blobs/{SPLIT}",
    scene_filter,
    sensor_config=SensorConfig.build_all_sensors(),
)

frame_idx = 1

# token = np.random.choice(scene_loader.tokens)
token = "ed4ac2dad0fa584b"
scene = scene_loader.get_scene_from_token(token)
fig, ax = plot_bev_frame(scene, frame_idx)
# plt.show()
plot_cameras_frame(scene, frame_idx)[0].savefig('./debug/ed4ac2dad0fa584b_vis_cam.png')
# fig.savefig('./debug/ed4ac2dad0fa584b_gt.png')