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

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

# %% auto 0
__all__ = ['BenchmarkSpecSimulation', 'BenchmarkSpecPrediction', 'TrainingContext', 'run_benchmark', 'run_multiple_benchmarks']

# %% ../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 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 run_multiple_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]|None = None, # Hyperparameters passed to build_model
    seed: int|None = None, # Base random seed
    continue_on_error: bool = True, # If True, continue running benchmarks even if one fails
) -> list[dict[str, Any]]:
    """
    Runs multiple benchmarks sequentially using the same build_model function.

    Args:
        specs: A list or dictionary containing the BenchmarkSpec objects to run.
        build_model: A callable that accepts a TrainingContext and returns a trained model/predictor function.
        hyperparameters: A dictionary of hyperparameters passed to the build_model function.
        seed: An optional integer seed passed to each run_benchmark call for reproducibility.
        continue_on_error: If True, catches exceptions during individual benchmark runs, prints a warning,
                           and continues. If False, stops on the first error.

    Returns:
        A list of result dictionaries containing the results from successful benchmark runs.
    """
    results_list = []
    hyperparameters = hyperparameters or {} # Ensure it's a dict

    # Determine the list of specification objects to iterate over
    spec_objects = list(specs.values()) if isinstance(specs, dict) else list(specs)

    print(f"--- Starting benchmark run for {len(spec_objects)} specifications ---")

    for i, spec in enumerate(spec_objects):
        spec_name = getattr(spec, 'name', f'Unnamed Spec {i+1}') # Get name for logging
        print(f"\n[{i+1}/{len(spec_objects)}] Running benchmark: {spec_name}")

        try:
            # Run the individual benchmark
            result = run_benchmark(
                spec=spec,
                build_model=build_model,
                hyperparameters=hyperparameters,
                seed=seed # Pass the same base seed to each run
            )
            results_list.append(result)
            print(f"  -> Success: {spec_name} completed.")

        except Exception as e:
            print(f"  -> ERROR running benchmark '{spec_name}': {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)}/{len(spec_objects)} completed successfully. ---")
    return results_list
