import argparse
from multiprocessing import Pool, set_start_method

import yaml

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


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="Legacy MEOW-ML runtime 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()
