import logging
import hydra
from omegaconf import DictConfig

from nuplan.planning.script.builders.logging_builder import build_logger

from navsim.planning.metric_caching.caching import cache_data
from navsim.planning.script.builders.worker_pool_builder import build_worker


logger = logging.getLogger(__name__)

CONFIG_PATH = "config/metric_caching"
# CONFIG_NAME = "cache_trainval_pt1"


# @hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME)
@hydra.main(config_path=CONFIG_PATH,)
def main(cfg: DictConfig) -> None:
    """
    Main entrypoint for training/validation experiments.
    :param cfg: omegaconf dictionary
    """
    # Configure logger
    build_logger(cfg)
    # Build worker
    worker = build_worker(cfg)

    # Precompute and cache all features
    logger.info("Starting Metric Caching...")
    if cfg.worker == "ray_distributed" and cfg.worker.use_distributed:
        raise AssertionError("ray in distributed mode will not work with this job")
    cache_data(cfg=cfg, worker=worker)

if __name__ == "__main__":
    main()