"""MEOW-ML main 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
from multiprocessing import Pool, set_start_method

import yaml

from .core.training.train import train
from .core.utils.mlflow_helpers import setup_mlflow
from meow_ml.src.hyperparam_optimization.tune import tune
from meow_ml.src.utils.gpu_setup import (
    init_cpu,
    split_runs_across_gpus,
    sub_process_entry_point,
)


def run(args):
    with open(args.config, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)
        if args.optuna:
            runs_config = config["studies"]
            data_generator_config = config["data_generator"]
            mlflow_config = config["mlflow"]
        else:
            runs_config = config["models"]
            data_generator_config = config["data_generator"]
            mlflow_config = config["mlflow"]

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

        if args.gpu:
            gpu_config = config["gpu"]

    if args.gpu:
        gpu_config["GPUS"] = [int(x) for x in gpu_config["GPUS"].split(",") if x.isdigit()]

        split_runs = split_runs_across_gpus(runs_config, len(gpu_config["GPUS"]))

        p1 = Pool(processes=len(gpu_config["GPUS"]))
        set_start_method("spawn", force=True)

        for gpu_id, run_config in zip(gpu_config["GPUS"], split_runs):
            p1.apply_async(
                sub_process_entry_point,
                args=(
                    run_config.tolist(),
                    data_generator_config,
                    mlflow_config,
                    gpu_id,
                    args.optuna,
                ),
            )

        p1.close()
        p1.join()
    else:
        init_cpu()
        if args.optuna:
            for study_variables in runs_config:
                tune(
                    study_variables=study_variables,
                    data_generator_config=data_generator_config,
                    mlflow_config=mlflow_config)
        else:
            for model_config in runs_config:
                train(
                    model_config=model_config,
                    data_generator_config=data_generator_config,
                    mlflow_config=mlflow_config,
                )


def build_parser():
    parser = argparse.ArgumentParser(description="MEOW-ML training entrypoint")
    parser.add_argument("--config", required=True, type=str, help="Path to a YAML config file")
    parser.add_argument("--gpu", action="store_true", help="Enable multi-GPU execution")
    parser.add_argument(
        "--optuna", action="store_true", help="Run hyperparameter optimization studies"
    )
    return parser


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


if __name__ == "__main__":
    main()
