"""MEOW-ML core training entrypoint.

Copyright ©2024 United States Government as represented by the
Administrator of the National Aeronautics and Space Administration.
No copyright is claimed in the United States under Title 17, U.S. Code.
All Other Rights Reserved.
"""

import argparse
import os

import yaml

from .training.train import train
from .utils.mlflow_helpers import setup_mlflow


def build_parser():
    parser = argparse.ArgumentParser(
        description="MEOW-ML: Multi-Path Earth Observation Machine Learning Framework"
    )
    parser.add_argument("--config", required=True, type=str, help="Path to config file")
    return parser


def main(argv=None):
    parser = build_parser()
    args = parser.parse_args(argv)

    with open(args.config, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)
        models_config = config.get("models", [])
        data_generator_config = config.get("data_generator", {})
        mlflow_config = config.get("mlflow", {})
        gpu_config = config.get("gpu", {})

    gpus = gpu_config.get("GPUS")
    if gpus:
        os.environ["CUDA_VISIBLE_DEVICES"] = gpus

    mlflow_config = setup_mlflow(mlflow_config, config_file_path=args.config)

    for model_config in models_config:
        train(
            model_config=model_config,
            data_generator_config=data_generator_config,
            mlflow_config=mlflow_config,
        )


if __name__ == "__main__":
    main()
