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

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

# %% auto 0
__all__ = ['BenchmarkWH_Simulation', 'BenchmarkWH_Prediction', 'BenchmarkSilverbox_Simulation', 'BenchmarkSilverbox_Prediction',
           'BenchmarkCascadedTanks_Simulation', 'BenchmarkCascadedTanks_Prediction', 'BenchmarkEMPS_Simulation',
           'BenchmarkEMPS_Prediction', 'BenchmarkNoisyWH_Simulation', 'BenchmarkNoisyWH_Prediction',
           'BenchmarkCED_Simulation', 'BenchmarkCED_Prediction', 'dl_wiener_hammerstein', 'dl_silverbox',
           'dl_cascaded_tanks', 'dl_emps', 'dl_noisy_wh', 'dl_ced']

# %% ../../nbs/datasets/workshop.ipynb 2
from ..utils import *
import identibench.benchmark as idb
import identibench.metrics
import nonlinear_benchmarks
import numpy as np
from nonlinear_benchmarks.utilities import Input_output_data
from pathlib import Path
import shutil

# %% ../../nbs/datasets/workshop.ipynb 7
def dl_wiener_hammerstein(
        save_path: Path, #directory the files are written to, created if it does not exist
        force_download: bool = False, # force download the dataset
        save_train_valid: bool = True, # save unsplitted train and valid datasets in 'train_valid' subdirectory
        split_idx: int = 80_000 # split index for train and valid datasets
) -> None:
    train_val, test = nonlinear_benchmarks.WienerHammerBenchMark(force_download=force_download)
    train = train_val[:split_idx]
    valid = train_val[split_idx:]

    dataset_to_hdf5(train,valid,test,save_path,train_valid=(train_val if save_train_valid else None))

# %% ../../nbs/datasets/workshop.ipynb 9
def rmse_mV(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    return identibench.metrics.rmse(y_true,y_pred)*1000

# %% ../../nbs/datasets/workshop.ipynb 10
BenchmarkWH_Simulation = idb.BenchmarkSpecSimulation(
    name='BenchmarkWH_Simulation', dataset_id='wh',
    u_cols=['u0'], y_cols=['y0'], 
    metric_func=rmse_mV, 
    download_func=dl_wiener_hammerstein,
    init_window=50
)
BenchmarkWH_Prediction = idb.BenchmarkSpecPrediction(
    name='BenchmarkWH_Prediction', dataset_id='wh',
    u_cols=['u0'], y_cols=['y0'], 
    metric_func=rmse_mV, 
    download_func=dl_wiener_hammerstein,
    init_window=50, pred_horizon=100, pred_step=100
)

# %% ../../nbs/datasets/workshop.ipynb 15
def dl_silverbox(
        save_path: Path, #directory the files are written to, created if it does not exist
        force_download: bool = False, # force download the dataset
        save_train_valid: bool = True, # save unsplitted train and valid datasets in 'train_valid' subdirectory
        split_idx: int = 50_000 # split index for train and valid datasets
) -> None:
    train_val, test = nonlinear_benchmarks.Silverbox(force_download=force_download)
    train = train_val[:split_idx]
    valid = train_val[split_idx:]

    dataset_to_hdf5(train,valid,test,save_path,train_valid=(train_val if save_train_valid else None))

# %% ../../nbs/datasets/workshop.ipynb 17
def evaluate_silverbox(results: list[tuple[np.ndarray, np.ndarray]], spec: idb.BenchmarkSpecBase) -> dict[str, float]:

    test_configs = [
        ('test_0.hdf5', 'multisine_rmse'),
        ('test_1.hdf5', 'arrow_full_rmse'),
        ('test_2.hdf5', 'arrow_no_extrapolation_rmse'),
    ]

    aggregated_scores = {}
    for filename_part, score_name in test_configs:
        # Find the index for the current filename part
        idx = next(i for i, s in enumerate(spec.test_files) if filename_part in s.name) 
        aggregated_scores.update(
            idb.aggregate_metric_score([results[idx]], spec.metric_func, score_name=score_name)
        )

    return aggregated_scores

# %% ../../nbs/datasets/workshop.ipynb 18
BenchmarkSilverbox_Simulation = idb.BenchmarkSpecSimulation(
    name='BenchmarkSilverbox_Simulation', dataset_id='silverbox',
    u_cols=['u0'], y_cols=['y0'], 
    metric_func=rmse_mV, 
    download_func=dl_silverbox,
    custom_test_evaluation=evaluate_silverbox,
    init_window=50
)
BenchmarkSilverbox_Prediction = idb.BenchmarkSpecPrediction(
    name='BenchmarkSilverbox_Prediction', dataset_id='silverbox',
    u_cols=['u0'], y_cols=['y0'], 
    metric_func=rmse_mV, 
    download_func=dl_silverbox,
    custom_test_evaluation=evaluate_silverbox,
    init_window=50, pred_horizon=100, pred_step=100
)

# %% ../../nbs/datasets/workshop.ipynb 23
def dl_cascaded_tanks(
        save_path: Path, #directory the files are written to, created if it does not exist
        force_download: bool = False, # force download the dataset
        save_train_valid: bool = True, # save unsplitted train and valid datasets in 'train_valid' subdirectory
        split_idx: int = 160 # split index for train and valid datasets
) -> None:
    train_val, test = nonlinear_benchmarks.Cascaded_Tanks(force_download=force_download)
    train = train_val[split_idx:]
    valid = train_val[:split_idx]

    dataset_to_hdf5(train,valid,test,save_path,train_valid=(train_val if save_train_valid else None))

# %% ../../nbs/datasets/workshop.ipynb 25
BenchmarkCascadedTanks_Simulation = idb.BenchmarkSpecSimulation(
    name='BenchmarkCascadedTanks_Simulation', dataset_id='cascaded_tanks',
    u_cols=['u0'], y_cols=['y0'], metric_func=identibench.metrics.rmse, 
    download_func=dl_cascaded_tanks,
    init_window=50
)
BenchmarkCascadedTanks_Prediction = idb.BenchmarkSpecPrediction(
    name='BenchmarkCascadedTanks_Prediction', dataset_id='cascaded_tanks',
    u_cols=['u0'], y_cols=['y0'], metric_func=identibench.metrics.rmse, 
    download_func=dl_cascaded_tanks,
    init_window=50, pred_horizon=100, pred_step=100
)

# %% ../../nbs/datasets/workshop.ipynb 30
def dl_emps(
        save_path: Path, #directory the files are written to, created if it does not exist
        force_download: bool = False, # force download the dataset
        save_train_valid: bool = True, # save unsplitted train and valid datasets in 'train_valid' subdirectory
        split_idx: int = 18_000 # split index for train and valid datasets
) -> None:
    train_val, test = nonlinear_benchmarks.EMPS(force_download=force_download)
    train = train_val[:split_idx]
    valid = train_val[split_idx:]

    dataset_to_hdf5(train,valid,test,save_path,train_valid=(train_val if save_train_valid else None))

# %% ../../nbs/datasets/workshop.ipynb 32
BenchmarkEMPS_Simulation = idb.BenchmarkSpecSimulation(
    name='BenchmarkEMPS_Simulation', dataset_id='emps',
    u_cols=['u0'], y_cols=['y0'], metric_func=rmse_mV, 
    download_func=dl_emps,
    init_window=20
)
BenchmarkEMPS_Prediction = idb.BenchmarkSpecPrediction(
    name='BenchmarkEMPS_Prediction', dataset_id='emps',
    u_cols=['u0'], y_cols=['y0'], metric_func=rmse_mV, 
    download_func=dl_emps,
    init_window=20, pred_horizon=500, pred_step=100
)

# %% ../../nbs/datasets/workshop.ipynb 36
from scipy.io import loadmat
def dl_noisy_wh(
        save_path: Path, #directory the files are written to, created if it does not exist
        force_download: bool = False, # force download the dataset
        save_train_valid: bool = True # save unsplitted train and valid datasets in 'train_valid' subdirectory
) -> None:
    'the wiener hammerstein dataset with process noise'

    #extract raw .mat files, to preserve filenames necessary for train, valid split
    matfiles = nonlinear_benchmarks.not_splitted_benchmarks.WienerHammerstein_Process_Noise(data_file_locations=True,train_test_split=False,force_download=force_download)

    for file in matfiles:
        f_path = Path(file)
        save_path = Path(save_path)

        if 'Test' in f_path.stem:
            hdf_path = save_path / 'test'
        elif 'Combined' in f_path.stem:
            hdf_path = save_path / 'valid'
        else:
            hdf_path = save_path / 'train'

        out = loadmat(f_path)
        _,u,y,fs = out['dataMeas'][0,0]
        fs = fs[0,0]
        for idx,(ui,yi) in enumerate(zip(u.T,y.T)):
            iodata = Input_output_data(u=ui,y=yi, sampling_time=1/fs)
            fname = f'{f_path.stem}_{idx+1}'
            iodata_to_hdf5(iodata,hdf_path,fname)
    if save_train_valid:
        #copy train and valid files to train_valid directory
        for d in ['train','valid']:
            for f in (Path(save_path)/d).glob('*.hdf5') :
                shutil.copy2(f, (p:=Path(save_path)/'train_valid').mkdir(exist_ok=True) or p)


# %% ../../nbs/datasets/workshop.ipynb 38
BenchmarkNoisyWH_Simulation = idb.BenchmarkSpecSimulation(
    name='BenchmarkNoisyWH_Simulation', dataset_id='noisy_wh',
    u_cols=['u0'], y_cols=['y0'], metric_func=rmse_mV, 
    download_func=dl_noisy_wh,
    init_window=100
)
BenchmarkNoisyWH_Prediction = idb.BenchmarkSpecPrediction(
    name='BenchmarkNoisyWH_Prediction', dataset_id='noisy_wh',
    u_cols=['u0'], y_cols=['y0'], metric_func=rmse_mV, 
    download_func=dl_noisy_wh,
    init_window=100, pred_horizon=100, pred_step=100
)

# %% ../../nbs/datasets/workshop.ipynb 47
def dl_ced(
        save_path: Path, #directory the files are written to, created if it does not exist
        force_download: bool = False, # force download the dataset
        save_train_valid: bool = True, # save unsplitted train and valid datasets in 'train_valid' subdirectory
        split_idx: int = 300 # split index for train and valid datasets
) -> None:
    train_val, test = nonlinear_benchmarks.CED(force_download=force_download,always_return_tuples_of_datasets=True)
    train = tuple(x[:split_idx] for x in train_val)
    valid = tuple(x[split_idx:] for x in train_val)

    dataset_to_hdf5(train,valid,test,save_path,train_valid=(train_val if save_train_valid else None))

# %% ../../nbs/datasets/workshop.ipynb 49
def evaluate_ced(results: list[tuple[np.ndarray, np.ndarray]], spec: idb.BenchmarkSpecBase) -> dict[str, float]:
    test_configs = [
        ('test_0.hdf5', 'test_1_rmse'),
        ('test_1.hdf5', 'test_2_rmse'),
    ]

    aggregated_scores = {}
    for filename_part, score_name in test_configs:
        # Find the index for the current filename part
        idx = next(i for i, s in enumerate(spec.test_files) if filename_part in s.name) 
        aggregated_scores.update(
            idb.aggregate_metric_score([results[idx]], spec.metric_func, score_name=score_name)
        )

    return aggregated_scores

# %% ../../nbs/datasets/workshop.ipynb 50
BenchmarkCED_Simulation = idb.BenchmarkSpecSimulation(
    name='BenchmarkCED_Simulation', dataset_id='ced',
    u_cols=['u0'], y_cols=['y0'], metric_func=identibench.metrics.rmse, 
    download_func=dl_ced,
    custom_test_evaluation=evaluate_ced,
    init_window=10
)
BenchmarkCED_Prediction = idb.BenchmarkSpecPrediction(
    name='BenchmarkCED_Prediction', dataset_id='ced',
    u_cols=['u0'], y_cols=['y0'], metric_func=identibench.metrics.rmse, 
    download_func=dl_ced,
    custom_test_evaluation=evaluate_ced,
    init_window=10, pred_horizon=30, pred_step=30
)
