"""Core benchmark specification classes and runner functions."""

__all__ = [
    "BenchmarkSpecSimulation",
    "BenchmarkSpecPrediction",
    "TrainingContext",
    "run_benchmark",
    "benchmark_results_to_dataframe",
    "run_benchmarks",
    "aggregate_benchmark_results",
]

from pathlib import Path
from collections.abc import Iterator, Callable
from typing import Any
import random
import multiprocessing
import numpy as np
import pandas as pd
import time
from .utils import get_default_data_root, _load_sequences_from_files, hdf_files_from_path


def aggregate_metric_score(
    test_results, metric_func, score_name=None, sequence_aggregation_func=np.mean, window_aggregation_func=np.mean
):
    # iterate over test_results and calculate metric score for each (y_pred,y_test) tuple, if prediction, iterate over nested tuples with nested loop
    if score_name is None:
        score_name = metric_func.__name__
    if isinstance(test_results[0], list):
        scores = []
        for windowed_sequence in test_results:
            scores.append(
                window_aggregation_func([metric_func(y_pred, y_test) for y_pred, y_test in windowed_sequence])
            )
    else:
        scores = [[metric_func(y_pred, y_test) for y_pred, y_test in test_results]]
    return {score_name: sequence_aggregation_func(scores)}


class BenchmarkSpecBase:
    pass


class BenchmarkSpecBase:
    """
    Base class for benchmark specifications, holding common attributes.
    """

    def __init__(
        self,
        name: str,  # Unique name identifying this benchmark task.
        dataset_id: str,  # Identifier for the raw dataset source.
        u_cols: list[str],  # list of column names for input signals (u).
        y_cols: list[str],  # list of column names for output signals (y).
        metric_func: Callable[[np.ndarray, np.ndarray], float],  # Primary metric: `func(y_true, y_pred)`.
        x_cols: list[str] | None = None,  # Optional state inputs (x).
        sampling_time: float | None = None,  # Optional sampling time (seconds).
        download_func: Callable[[Path, bool], None] | None = None,  # Dataset preparation func.
        test_model_func: Callable[[BenchmarkSpecBase, Callable], dict[str, Any]] = None,
        custom_test_evaluation=None,
        init_window: int | None = None,  # Steps for warm-up, potentially ignored in evaluation.
        data_root: [
            Path,
            Callable[[], Path],
        ] = get_default_data_root,  # root dir for dataset, may be a callable or path
    ):
        self.name = name
        self.dataset_id = dataset_id
        self.u_cols = u_cols
        self.y_cols = y_cols
        self.metric_func = metric_func
        self.x_cols = x_cols
        self.sampling_time = sampling_time
        self.download_func = download_func
        self.test_model_func = test_model_func
        self.custom_test_evaluation = custom_test_evaluation
        self.init_window = init_window
        self._data_root = data_root

    @property
    def data_root(self) -> Path:
        """Returns the evaluated data root path."""
        if isinstance(self._data_root, Callable):
            return self._data_root()
        return self._data_root

    @property
    def dataset_path(self) -> Path:
        """Returns the full path to the dataset directory."""
        return self.data_root / self.dataset_id

    @property
    def test_path(self) -> Path:
        """Returns the full path to the directory containing the test files ."""
        return self.dataset_path / "test"

    @property
    def train_path(self) -> Path:
        """Returns the full path to the directory containing the train files."""
        return self.dataset_path / "train"

    @property
    def valid_path(self) -> Path:
        """Returns the full path to the directory containing the valid files."""
        return self.dataset_path / "valid"

    @property
    def train_valid_path(self) -> Path:
        """Returns the full path to the directory containing the train_valid files."""
        return self.dataset_path / "train_valid"

    @property
    def train_files(self) -> list[Path]:
        """Returns the list of hdf5 files in the training directory."""
        return hdf_files_from_path(self.train_path)

    @property
    def valid_files(self) -> list[Path]:
        """Returns the list of hdf5 files in the validation directory."""
        return hdf_files_from_path(self.valid_path)

    @property
    def train_valid_files(self) -> list[Path]:
        """Returns the list of hdf5 files in the train_valid directory if it exists, otherwise returns the union of the train and valid directories."""
        train_valid_files = hdf_files_from_path(self.train_valid_path)
        if train_valid_files:
            return train_valid_files
        else:
            train_files = hdf_files_from_path(self.train_path)
            valid_files = hdf_files_from_path(self.valid_path)
            return sorted(train_files + valid_files)

    @property
    def test_files(self) -> list[Path]:
        """Returns the list of hdf5 files in the test directory."""
        return hdf_files_from_path(self.test_path)

    def ensure_dataset_exists(self, force_download: bool = False) -> None:
        """Checks if the dataset exists, downloads/prepares it if needed.

        Downloads are run in a subprocess to isolate side effects from problematic
        dependencies (e.g., nest_asyncio, rospy) that can corrupt the parent process.
        """
        dataset_path = self.dataset_path
        if self.download_func is None:
            print(f"Warning: No download function for '{self.name}'. Assuming data exists at {dataset_path}")
            if not dataset_path.is_dir():
                print(f"Warning: Dataset directory {dataset_path} not found.")
            return

        if not dataset_path.is_dir() or force_download:
            print(f"Preparing dataset for '{self.name}' at {dataset_path}...")
            self.data_root.mkdir(parents=True, exist_ok=True)
            ctx = multiprocessing.get_context("spawn")
            p = ctx.Process(
                target=self.download_func,
                args=(dataset_path, force_download),
            )
            p.start()
            p.join()
            if p.exitcode != 0:
                raise RuntimeError(f"Download for '{self.name}' failed (exit code {p.exitcode})")
            print(f"Dataset '{self.name}' prepared successfully.")


def _test_simulation(specs, model):
    results = []
    for u_test, y_test, _ in _load_sequences_from_files(specs.test_files, specs.u_cols, specs.y_cols, specs.x_cols):
        y_pred = model(u_test, y_test[: specs.init_window])
        y_test_win = y_test[specs.init_window :]
        y_pred = y_pred[-y_test_win.shape[0] :]
        results.append((y_pred, y_test_win))
    return results


class BenchmarkSpecSimulation(BenchmarkSpecBase):
    """
    Specification for a simulation benchmark task.

    Inherits common parameters from BaseBenchmarkSpec.
    Use this when the goal is to simulate the system's output given the input `u`.
    """

    def __init__(
        self,
        name: str,  # Unique name identifying this benchmark task.
        dataset_id: str,  # Identifier for the raw dataset source.
        u_cols: list[str],  # list of column names for input signals (u).
        y_cols: list[str],  # list of column names for output signals (y).
        metric_func: Callable[[np.ndarray, np.ndarray], float],  # Primary metric: `func(y_true, y_pred)`.
        x_cols: list[str] | None = None,  # Optional state inputs (x).
        sampling_time: float | None = None,  # Optional sampling time (seconds).
        download_func: Callable[[Path, bool], None] | None = None,  # Dataset preparation func.
        test_model_func: Callable[[BenchmarkSpecBase, Callable], dict[str, Any]] = _test_simulation,
        custom_test_evaluation=None,
        init_window: int | None = None,  # Steps for warm-up, potentially ignored in evaluation.
        data_root: [
            Path,
            Callable[[], Path],
        ] = get_default_data_root,  # root dir for dataset, may be a callable or path
    ):
        self.name = name
        self.dataset_id = dataset_id
        self.u_cols = u_cols
        self.y_cols = y_cols
        self.metric_func = metric_func
        self.x_cols = x_cols
        self.sampling_time = sampling_time
        self.download_func = download_func
        self.test_model_func = test_model_func
        self.custom_test_evaluation = custom_test_evaluation
        self.init_window = init_window
        self._data_root = data_root


def _test_prediction(specs, model):
    results = []
    for u_test, y_test, _ in _load_sequences_from_files(specs.test_files, specs.u_cols, specs.y_cols, specs.x_cols):
        # iterate through windows of u_test and y_test
        window_results = []
        for i in range(0, u_test.shape[0] - specs.init_window - specs.pred_horizon, specs.pred_step):
            u_test_win = u_test[i : i + specs.init_window + specs.pred_horizon]
            y_test_win = y_test[i : i + specs.init_window + specs.pred_horizon]
            y_pred = model(u_test_win, y_test_win[: specs.init_window])
            window_results.append((y_pred, y_test_win))
        results.append(window_results)
    return results


class BenchmarkSpecPrediction(BenchmarkSpecBase):
    """
    Specification for a k-step ahead prediction benchmark task.

    Inherits common parameters from BaseBenchmarkSpec and adds prediction-specific ones.
    Use this when the goal is to predict `y` some steps ahead based on past `u` and `y`.
    """

    def __init__(
        self,
        name: str,  # Unique name identifying this benchmark task.
        dataset_id: str,  # Identifier for the raw dataset source.
        u_cols: list[str],  # list of column names for input signals (u).
        y_cols: list[str],  # list of column names for output signals (y).
        metric_func: Callable[[np.ndarray, np.ndarray], float],  # Primary metric: `func(y_true, y_pred)`.
        pred_horizon: int,  # The 'k' in k-step ahead prediction (mandatory for this type).
        pred_step: int,  # Step size for k-step ahead prediction (e.g., predict y[t+k] using data up to t).
        x_cols: list[str] | None = None,  # Optional state inputs (x).
        sampling_time: float | None = None,  # Optional sampling time (seconds).
        download_func: Callable[[Path, bool], None] | None = None,  # Dataset preparation func.
        test_model_func: Callable[[BenchmarkSpecBase, Callable], dict[str, Any]] = _test_prediction,
        custom_test_evaluation=None,
        init_window: int | None = None,  # Steps for warm-up, potentially ignored in evaluation.
        data_root: [
            Path,
            Callable[[], Path],
        ] = get_default_data_root,  # root dir for dataset, may be a callable or path
    ):
        if pred_horizon <= 0:
            raise ValueError("pred_horizon must be a positive integer for PredictionBenchmarkSpec.")
        self.pred_horizon = pred_horizon
        self.pred_step = pred_step

        self.name = name
        self.dataset_id = dataset_id
        self.u_cols = u_cols
        self.y_cols = y_cols
        self.metric_func = metric_func
        self.x_cols = x_cols
        self.sampling_time = sampling_time
        self.download_func = download_func
        self.test_model_func = test_model_func
        self.custom_test_evaluation = custom_test_evaluation
        self.init_window = init_window
        self._data_root = data_root


class TrainingContext:
    """
    Context object passed to the user's training function (`build_predictor`).

    Holds the benchmark specification, hyperparameters, and seed.
    Provides methods to access the raw, full-length training and validation data sequences.
    Windowing/batching for training must be handled within the user's `build_predictor` function.
    """

    def __init__(
        self,
        spec: BenchmarkSpecBase,  # The benchmark specification.
        hyperparameters: dict[str, Any],  # User-provided dictionary containing model and training hyperparameters.
        seed: int | None = None,  # Optional random seed for reproducibility.
    ):
        self.spec = spec
        self.hyperparameters = hyperparameters
        self.seed = seed

    # --- Data Access Methods ---

    def get_train_sequences(self) -> Iterator[tuple[np.ndarray, np.ndarray, np.ndarray | None]]:
        """Returns a lazy iterator yielding raw (u, y, x) tuples for the 'train' subset."""
        return _load_sequences_from_files(
            file_paths=self.spec.train_files,
            u_cols=self.spec.u_cols,
            y_cols=self.spec.y_cols,
            x_cols=self.spec.x_cols,
        )

    def get_valid_sequences(self) -> Iterator[tuple[np.ndarray, np.ndarray, np.ndarray | None]]:
        """Returns a lazy iterator yielding raw (u, y, x) tuples for the 'valid' subset."""
        return _load_sequences_from_files(
            file_paths=self.spec.valid_files,
            u_cols=self.spec.u_cols,
            y_cols=self.spec.y_cols,
            x_cols=self.spec.x_cols,
        )

    def get_train_valid_sequences(self) -> Iterator[tuple[np.ndarray, np.ndarray, np.ndarray | None]]:
        """
        Returns a lazy iterator yielding raw (u, y, x) tuples for combined training and validation.

        Checks for a 'train_valid' subset directory first. If it exists, loads data from there.
        If not, it loads data from 'train' and 'valid' subsets sequentially.
        """
        return _load_sequences_from_files(
            file_paths=self.spec.train_valid_files,
            u_cols=self.spec.u_cols,
            y_cols=self.spec.y_cols,
            x_cols=self.spec.x_cols,
        )


def run_benchmark(spec, build_model, hyperparameters={}, seed=None):

    if seed is None:
        seed = random.randint(0, 2**32 - 1)

    results = {
        "benchmark_name": spec.name,
        "dataset_id": spec.dataset_id,
        "hyperparameters": hyperparameters,
        "seed": seed,
        "training_time_seconds": np.nan,
        "test_time_seconds": np.nan,
        "benchmark_type": type(spec).__name__,
        "metric_name": spec.metric_func.__name__,
        "metric_score": np.nan,
        "custom_scores": {},
        "model_predictions": [],
    }

    spec.ensure_dataset_exists()

    context = TrainingContext(spec=spec, hyperparameters=hyperparameters, seed=seed)

    train_start_time = time.monotonic()
    model = build_model(context)
    train_end_time = time.monotonic()
    results["training_time_seconds"] = train_end_time - train_start_time

    if model is None:
        raise RuntimeError(f"build_model for {spec.name} did not return a model.")

    test_start_time = time.monotonic()
    test_results = spec.test_model_func(spec, model)  # get list of (y_pred,y_test) tuples
    test_end_time = time.monotonic()
    results["test_time_seconds"] = test_end_time - test_start_time

    results["model_predictions"] = test_results

    results.update(aggregate_metric_score(test_results, spec.metric_func, score_name="metric_score"))
    if spec.custom_test_evaluation is not None:
        results["custom_scores"] = spec.custom_test_evaluation(test_results, spec)

    return results


# Define a very simple build_model function for the example
def _dummy_build_model(context):
    print(f"Building model with spec: {context.spec.name}, seed: {context.seed}")

    def dummy_model(u_test, y_test):
        output_dim = len(context.spec.y_cols)
        return np.zeros((u_test.shape[0], output_dim))

    return dummy_model  # Return the callable model


def benchmark_results_to_dataframe(
    results_list: list[dict[str, Any]],  # List of benchmark result dictionaries from `run_benchmark`.
) -> pd.DataFrame:
    """Transforms a list of benchmark result dictionaries into a pandas DataFrame."""
    if not results_list:
        return pd.DataFrame()

    df = pd.DataFrame(results_list)

    # Flatten custom_scores if the column exists and contains dictionaries
    if "custom_scores" in df.columns and df["custom_scores"].apply(isinstance, args=(dict,)).any():
        # Fill missing/non-dict entries with empty dicts for normalization
        custom_scores_filled = df["custom_scores"].apply(lambda x: x if isinstance(x, dict) else {})
        custom_scores_df = pd.json_normalize(custom_scores_filled).add_prefix("cs_")
        # Drop the original nested column and join the flattened one
        df = pd.concat([df.drop(columns=["custom_scores"]), custom_scores_df], axis=1)

    # Drop the 'model_predictions' column as it's usually large and not suitable for a summary DataFrame
    if "model_predictions" in df.columns:
        df = df.drop(columns=["model_predictions"])

    return df


def run_benchmarks(
    specs: list[BenchmarkSpecBase] | dict[str, BenchmarkSpecBase],  # Collection of specs to run.
    build_model: Callable[[TrainingContext], Callable],  # User function to build the model/predictor.
    hyperparameters: dict[str, Any]
    | list[dict[str, Any]]
    | None = None,  # Single dict, list of dicts (matching specs), or None.
    n_times: int = 1,  # Number of times to repeat each benchmark specification.
    continue_on_error: bool = True,  # If True, continue running benchmarks even if one fails.
    return_dataframe: bool = True,  # If True, return results as a pandas DataFrame, otherwise return a list of dicts.
) -> pd.DataFrame | list[dict[str, Any]]:
    """
    Runs multiple benchmarks sequentially, with repetitions and flexible hyperparameters.

    Returns either a pandas DataFrame summarizing the results (default)
    or a list of raw result dictionaries.
    """
    results_list = []
    spec_objects = list(specs.values()) if isinstance(specs, dict) else list(specs)
    num_specs = len(spec_objects)

    # Validate hyperparameters input
    if isinstance(hyperparameters, list):
        if len(hyperparameters) != num_specs:
            raise ValueError(
                f"If hyperparameters is a list, its length ({len(hyperparameters)}) must match the number of specs ({num_specs})."
            )
        get_hps = lambda i: hyperparameters[i]  # Function to get hp based on spec index
    else:
        # If None or a single dict, use the same for all. Ensure it's a dict or None.
        hps_single = hyperparameters or {}
        get_hps = lambda i: hps_single  # Function always returns the same hp dict

    print(f"--- Starting benchmark run for {num_specs} specifications, repeating each {n_times} times ---")

    total_runs = num_specs * n_times
    current_run = 0

    for repetition in range(n_times):
        print(f"\n-- Repetition {repetition + 1}/{n_times} --")
        for i, spec in enumerate(spec_objects):
            current_run += 1
            spec_name = getattr(spec, "name", f"Unnamed Spec {i + 1}")
            print(f"\n[{current_run}/{total_runs}] Running: {spec_name} (Rep {repetition + 1})")

            current_hyperparameters = get_hps(i)

            try:
                result = run_benchmark(spec=spec, build_model=build_model, hyperparameters=current_hyperparameters)
                results_list.append(result)
                print(f"  -> Success: {spec_name} (Rep {repetition + 1}) completed.")

            except Exception as e:
                print(f"  -> ERROR running benchmark '{spec_name}' (Rep {repetition + 1}): {e}")
                if not continue_on_error:
                    print("Stopping due to error (continue_on_error=False).")
                    raise

    print(f"\n--- Benchmark run finished. {len(results_list)}/{total_runs} individual runs completed successfully. ---")

    if return_dataframe:
        return benchmark_results_to_dataframe(results_list)
    else:
        return results_list


def aggregate_benchmark_results(
    results_df: pd.DataFrame,  # DataFrame returned by run_benchmarks (with return_dataframe=True).
    group_by_cols: str | list[str] = "benchmark_name",  # Column(s) to group by before aggregation.
    agg_funcs: str | list[str] = "mean",  # Aggregation function(s) ('mean', 'median', 'std', etc.) or list thereof.
) -> pd.DataFrame:
    """
    Aggregates numeric results from a benchmark DataFrame, grouped by specified columns.
    """
    if results_df.empty:
        return pd.DataFrame()  # Return empty if input is empty

    # Identify numeric columns suitable for aggregation
    numeric_cols = results_df.select_dtypes(include=np.number).columns.tolist()

    # Exclude columns that are typically identifiers or settings, even if numeric
    cols_to_exclude = ["seed", "repetition"]
    agg_cols = [
        col
        for col in numeric_cols
        if col not in cols_to_exclude
        and col not in (group_by_cols if isinstance(group_by_cols, list) else [group_by_cols])
    ]

    if not agg_cols:
        print("Warning: No numeric columns found to aggregate (excluding identifiers). Returning empty DataFrame.")
        return pd.DataFrame()

    try:
        # Perform groupby and aggregation
        aggregated_df = results_df.groupby(group_by_cols)[agg_cols].agg(agg_funcs)
    except Exception as e:
        print(f"Error during aggregation: {e}")
        # Provide more context if grouping fails
        if isinstance(group_by_cols, list):
            missing_cols = [col for col in group_by_cols if col not in results_df.columns]
        else:
            missing_cols = [group_by_cols] if group_by_cols not in results_df.columns else []

        if missing_cols:
            print(f"  -> Grouping columns not found in DataFrame: {missing_cols}")
        print(f"  -> Available columns: {results_df.columns.tolist()}")
        print(f"  -> Columns selected for aggregation: {agg_cols}")
        return pd.DataFrame()  # Return empty on error

    return aggregated_df
