#!/usr/bin/env python3

"""CLI module for running raytune tuning experiment."""

import argparse

import logging

import shutil

from pathlib import Path

from typing import Any

import ray

import yaml

from stimulus.data import loaders

from stimulus.learner import raytune_learner, raytune_parser

from stimulus.utils import launch_utils, yaml_data, yaml_model_schema

logger = logging.getLogger(__name__)

def _raise_empty_grid() -> None:

    """Raise an error when grid results are empty."""

    raise RuntimeError("Ray Tune returned empty results grid")

def get_args() -> argparse.Namespace:

    """Get the arguments when using from the commandline.

    Returns:

        Parsed command line arguments.

    """

    parser = argparse.ArgumentParser(description="Launch check_model.")

    parser.add_argument("-d", "--data", type=str, required=True, metavar="FILE", help="Path to input csv file.")

    parser.add_argument("-m", "--model", type=str, required=True, metavar="FILE", help="Path to model file.")

    parser.add_argument(

        "-e",

        "--data_config",

        type=str,

        required=True,

        metavar="FILE",

        help="Path to data config file.",

    )

    parser.add_argument(

        "-c",

        "--model_config",

        type=str,

        required=True,

        metavar="FILE",

        help="Path to yaml config training file.",

    )

    parser.add_argument(

        "-w",

        "--initial_weights",

        type=str,

        required=False,

        nargs="?",

        const=None,

        default=None,

        metavar="FILE",

        help="The path to the initial weights (optional).",

    )

    parser.add_argument(

        "--ray_results_dirpath",

        type=str,

        required=False,

        nargs="?",

        const=None,

        default=None,

        metavar="DIR_PATH",

        help="Location where ray_results output dir should be written. If None, uses ~/ray_results.",

    )

    parser.add_argument(

        "-o",

        "--output",

        type=str,

        required=False,

        nargs="?",

        const="best_model.pt",

        default="best_model.pt",

        metavar="FILE",

        help="The output file path to write the trained model to",

    )

    parser.add_argument(

        "-bm",

        "--best_metrics",

        type=str,

        required=False,

        nargs="?",

        const="best_metrics.csv",

        default="best_metrics.csv",

        metavar="FILE",

        help="The path to write the best metrics to",

    )

    parser.add_argument(

        "-bc",

        "--best_config",

        type=str,

        required=False,

        nargs="?",

        const="best_config.yaml",

        default="best_config.yaml",

        metavar="FILE",

        help="The path to write the best config to",

    )

    parser.add_argument(

        "-bo",

        "--best_optimizer",

        type=str,

        required=False,

        nargs="?",

        const="best_optimizer.pt",

        default="best_optimizer.pt",

        metavar="FILE",

        help="The path to write the best optimizer to",

    )

    parser.add_argument(

        "--tune_run_name",

        type=str,

        required=False,

        nargs="?",

        const=None,

        default=None,

        metavar="CUSTOM_RUN_NAME",

        help=(

            "Tells ray tune what the 'experiment_name' (i.e. the given tune_run name) should be. "

            "If set, the subdirectory of ray_results is named with this value and its train dir is prefixed accordingly. "

            "Default None means that ray will generate such a name on its own."

        ),

    )

    parser.add_argument(

        "--debug_mode",

        action="store_true",

        help="Activate debug mode for tuning. Default false, no debug.",

    )

    return parser.parse_args()

def main(

    model_path: str,

    data_path: str,

    data_config_path: str,

    model_config_path: str,

    initial_weights: str | None = None,  # noqa: ARG001

    ray_results_dirpath: str | None = None,

    output_path: str | None = None,

    best_optimizer_path: str | None = None,

    best_metrics_path: str | None = None,

    best_config_path: str | None = None,

    *,

    debug_mode: bool = False,

) -> None:

    """Run the main model checking pipeline.

    Args:

        data_path: Path to input data file.

        model_path: Path to model file.

        data_config_path: Path to data config file.

        model_config_path: Path to model config file.

        initial_weights: Optional path to initial weights.

        ray_results_dirpath: Directory for ray results.

        debug_mode: Whether to run in debug mode.

        output_path: Path to write the best model to.

        best_optimizer_path: Path to write the best optimizer to.

        best_metrics_path: Path to write the best metrics to.

        best_config_path: Path to write the best config to.

    """

    # Convert data config to proper type

    with open(data_config_path) as file:

        data_config_dict: dict[str, Any] = yaml.safe_load(file)

    data_config: yaml_data.YamlSubConfigDict = yaml_data.YamlSubConfigDict(**data_config_dict)

    with open(model_config_path) as file:

        model_config_dict: dict[str, Any] = yaml.safe_load(file)

    model_config: yaml_model_schema.Model = yaml_model_schema.Model(**model_config_dict)

    encoder_loader = loaders.EncoderLoader()

    encoder_loader.initialize_column_encoders_from_config(column_config=data_config.columns)

    model_class = launch_utils.import_class_from_file(model_path)

    ray_config_loader = yaml_model_schema.YamlRayConfigLoader(model=model_config)

    ray_config_model = ray_config_loader.get_config()

    tuner = raytune_learner.TuneWrapper(

        model_config=ray_config_model,

        data_config_path=data_config_path,

        model_class=model_class,

        data_path=data_path,

        encoder_loader=encoder_loader,

        seed=42,

        ray_results_dir=ray_results_dirpath,

        debug=debug_mode,

    )

    # Ensure output_path is provided

    if output_path is None:

        raise ValueError("output_path must not be None")

    try:

        grid_results = tuner.tune()

        if not grid_results:

            _raise_empty_grid()

        # Initialize parser with results

        parser = raytune_parser.TuneParser(result=grid_results)

        # Ensure output directory exists

        Path(output_path).parent.mkdir(parents=True, exist_ok=True)

        # Save outputs using proper Result object API

        parser.save_best_model(output=output_path)

        parser.save_best_optimizer(output=best_optimizer_path)

        parser.save_best_metrics_dataframe(output=best_metrics_path)

        parser.save_best_config(output=best_config_path)

    except RuntimeError:

        logger.exception("Tuning failed")

        raise

    except KeyError:

        logger.exception("Missing expected result key")

        raise

    finally:

        if debug_mode:

            logger.info("Debug mode - preserving Ray results directory")

        #elif ray_results_dirpath:

        #    shutil.rmtree(ray_results_dirpath, ignore_errors=True)

def run() -> None:

    """Run the model checking script."""

    args = get_args()

    main(

        data_path=args.data,

        model_path=args.model,

        data_config_path=args.data_config,

        model_config_path=args.model_config,

        initial_weights=args.initial_weights,

        ray_results_dirpath=args.ray_results_dirpath,

        output_path=args.output,

        best_optimizer_path=args.best_optimizer,

        best_metrics_path=args.best_metrics,

        best_config_path=args.best_config,

        debug_mode=args.debug_mode,

    )

if __name__ == "__main__":

    ray.init(address="auto", ignore_reinit_error=True)

    run()