"""Forest canopy example entrypoint for MEOW-ML.

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
from pathlib import Path

import yaml

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


DEFAULT_CONFIG = Path(__file__).with_suffix("").parent / "configs" / "chc_config.yaml"


def build_parser():
    parser = argparse.ArgumentParser(
        description="Forest canopy example using the MEOW-ML framework"
    )
    parser.add_argument(
        "--config",
        default=str(DEFAULT_CONFIG),
        type=str,
        help="Path to the canopy-height-change example 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
        print(f"Using GPUs: {gpus}")

    mlflow_config = setup_mlflow(mlflow_config, config_file_path=args.config)
    print(f"MLflow experiment: {mlflow_config.get('EXPERIMENT_NAME')}")

    for model_config in models_config:
        print(f"\nStarting training for model: {model_config.get('model_name', 'Unknown')}")
        train(
            model_config=model_config,
            data_generator_config=data_generator_config,
            mlflow_config=mlflow_config,
        )
        print(f"Completed training for model: {model_config.get('model_name', 'Unknown')}")


if __name__ == "__main__":
    main()
