"""Fill in a module description here"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/benchmark.ipynb.

# %% auto 0
__all__ = ['BenchmarkSpecSimulation', 'BenchmarkSpecPrediction', 'TrainingContext', 'run_benchmark',
           'benchmark_results_to_dataframe', 'run_benchmarks', 'aggregate_benchmark_results']

# %% ../nbs/benchmark.ipynb 2
from pathlib import Path
from collections.abc import Iterator, Callable
from typing import Any
import random
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

# %% ../nbs/benchmark.ipynb 5
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)}

# %% ../nbs/benchmark.ipynb 6
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."""
        # (Implementation remains the same as before)
        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)
            try:
                self.download_func(dataset_path, force_download)
                print(f"Dataset '{self.name}' prepared successfully.")
            except Exception as e:
                print(f"Error preparing dataset '{self.name}': {e}")
                raise

# %% ../nbs/benchmark.ipynb 7
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

# %% ../nbs/benchmark.ipynb 8
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

# %% ../nbs/benchmark.ipynb 9
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

# %% ../nbs/benchmark.ipynb 10
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

# %% ../nbs/benchmark.ipynb 17
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.
    """
    # Explicit __init__ for nbdev documentation compatibility
    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.
                ):
        # Standard attribute assignment
        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,
        )

# %% ../nbs/benchmark.ipynb 20
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

# %% ../nbs/benchmark.ipynb 21
# 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

# %% ../nbs/benchmark.ipynb 30
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

# %% ../nbs/benchmark.ipynb 31
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
                )
                # result['repetition'] = repetition # Add repetition number to results
                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

# %% ../nbs/benchmark.ipynb 34
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
