"""Entry point: load configs and run training.

Usage: (commands for each preset)
    python -m src.main gtauav_balanced
    python -m src.main gtauav_balanced_asym
    python -m src.main gtauav_balanced_stripnet
    python -m src.main gtauav_balanced_stripnet_unfrozen
    python -m src.main gtauav_baseline
    python -m src.main gtauav_baseline_asym
    python -m src.main gtauav_baseline_stripnet
    python -m src.main gtauav_baseline_stripnet_unfrozen
    python -m src.main gtauav_image_heavy
    python -m src.main gtauav_gtauav_text_heavy
"""

from __future__ import annotations

import logging
import sys

import coloredlogs

from src.conf.config_loader import load_all_configs
from src.training.trainer_new import Trainer
from src.utils.path_utils import get_proj_dir

logger = logging.getLogger("caption_test")


def main() -> None:
    coloredlogs.install(
        level="INFO",
        logger=logger,
        fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
    )

    if len(sys.argv) != 2:
        raise SystemExit(
            "Usage: python -m src.main <preset_name>\n"
            "Example: python -m src.main gtauav_balanced\n"
            "  available presets are subdirectories under in/config_files/",
        )
    preset_name = sys.argv[1]

    proj_dir = get_proj_dir()
    path2cfg = f"{proj_dir}in/config_files/"           # per REQUIREMENTS_GIN_STYLE.md §5

    # -------------------------------------------------------
    ''' ONLY FOR DEBUG with launch.json config:
         "args": ["main gtauav_balanced"] -> so need to extract
         preset name "gtauav_balanced"
    '''
    preset_name = preset_name.split(' ')[1]
    # -------------------------------------------------------

    configs = load_all_configs(path2cfg, preset_name)

    trainer = Trainer(
        pipeline_cfg=configs["pipeline"],
        hardware_cfg=configs["hardware"],
        training_cfg=configs["training"],
        tracking_cfg=configs["tracking"],
        models_common_cfg=configs["models_common"],
        models_cfg=configs["models"],
    )

    trainer.train()


if __name__ == "__main__":
    main()