"""Workshop benchmark dataset definitions (WH, Silverbox, Tanks, EMPS, NoisyWH, CED)."""

__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",
]

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


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))


def rmse_mV(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    return identibench.metrics.rmse(y_true, y_pred) * 1000


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,
)


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))


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


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,
)


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))


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,
)


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))


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,
)


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)


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,
)


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))


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


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,
)
