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

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

# %% auto 0
__all__ = ['u_forward', 'y_forward', 'BenchmarkRobotForward_Simulation', 'BenchmarkRobotForward_Prediction', 'u_inverse',
           'y_inverse', 'BenchmarkRobotInverse_Simulation', 'BenchmarkRobotInverse_Prediction', 'dl_robot_forward',
           'dl_robot_inverse']

# %% ../../nbs/datasets/industrial_robot.ipynb 2
from ..utils import write_dataset,write_array
import identibench.benchmark as idb
import identibench.metrics
from nonlinear_benchmarks.utilities import cashed_download
from pathlib import Path
import os
import h5py
import numpy as np
import scipy.io as sio

# %% ../../nbs/datasets/industrial_robot.ipynb 3
def robot_mat2hdf(
        save_path: Path, #directory the files are written to, created if it does not exist
        mat_path: Path, #path of mat file to extract
) -> None:         
    'converts .mat file of industrial robot to hdf5 file, used for forward and inverse'


    fs = 10  # Hz
    train_valid_split = 0.8
    os.makedirs(save_path / 'test', exist_ok=True)
    os.makedirs(save_path / 'train', exist_ok=True)
    os.makedirs(save_path / 'valid', exist_ok=True)

    mf = sio.loadmat(mat_path)
    for mode in ['train', 'test']:
        if mode == 'test':
            with h5py.File(save_path / 'test' / 'test.hdf5', 'w') as f:
                write_dataset(f, 'dt', np.ones_like(mf[f'time_{mode}'][0]) / fs)
                write_array(f, 'u', mf[f'u_{mode}'].T)
                write_array(f, 'y', mf[f'y_{mode}'].T)
                f.attrs['fs'] = fs

        else:
            with h5py.File(save_path / 'train' / 'train.hdf5', 'w') as train_f, \
                h5py.File(save_path / 'valid' / 'valid.hdf5', 'w') as valid_f:
                    dt = np.ones_like(mf[f'time_{mode}'][0]) / fs
                    total_entries = len(dt)
                    split_index = int(total_entries * train_valid_split)

                    write_dataset(train_f, 'dt', dt[:split_index])
                    write_array(train_f, 'u', mf[f'u_{mode}'][:,:split_index].T)
                    write_array(train_f, 'y', mf[f'y_{mode}'][:,:split_index].T)
                    train_f.attrs['fs'] = fs  
                    
                    write_dataset(valid_f, 'dt', dt[split_index:])
                    write_array(valid_f, 'u', mf[f'u_{mode}'][:,split_index:].T)
                    write_array(valid_f, 'y', mf[f'y_{mode}'][:,split_index:].T)
                    valid_f.attrs['fs'] = fs  
    

# %% ../../nbs/datasets/industrial_robot.ipynb 4
def dl_robot_forward(
        save_path: Path, #directory the files are written to, created if it does not exist
        force_download: bool = False # force download the dataset
) -> None:
    save_path = Path(save_path)
    url_robot = "https://fdm-fallback.uni-kl.de/TUK/FB/MV/WSKL/0001/Robot_Identification_Benchmark_Without_Raw_Data.rar"
    # unrar_download(url_robot,tmp_dir)

    tmp_dir = cashed_download(url_robot,'Industrial_robot',force_download=force_download)
    tmp_dir = Path(tmp_dir)

    path_forward = tmp_dir / "forward_identification_without_raw_data.mat"

    robot_mat2hdf(save_path,path_forward)
    

# %% ../../nbs/datasets/industrial_robot.ipynb 6
u_forward = [ f'u{i}' for i in range(0,6) ]
y_forward = [ f'y{i}' for i in range(0,6) ]

BenchmarkRobotForward_Simulation = idb.BenchmarkSpecSimulation(
    name='BenchmarkRobotForward_Simulation', dataset_id='robot_forward',
    u_cols=u_forward, y_cols=y_forward, metric_func=identibench.metrics.rmse, 
    download_func=dl_robot_forward,
    init_window=100
)
BenchmarkRobotForward_Prediction = idb.BenchmarkSpecPrediction(
    name='BenchmarkRobotForward_Prediction', dataset_id='robot_forward',
    u_cols=u_forward, y_cols=y_forward, metric_func=identibench.metrics.rmse, 
    download_func=dl_robot_forward,
    init_window=100, pred_horizon=100,pred_step=100
)

# %% ../../nbs/datasets/industrial_robot.ipynb 9
def dl_robot_inverse(
        save_path: Path, #directory the files are written to, created if it does not exist
        force_download: bool = False # force download the dataset
) -> None:
    save_path = Path(save_path)
    url_robot = "https://fdm-fallback.uni-kl.de/TUK/FB/MV/WSKL/0001/Robot_Identification_Benchmark_Without_Raw_Data.rar"
    # unrar_download(url_robot,tmp_dir)

    tmp_dir = cashed_download(url_robot,'Industrial_robot',force_download=force_download)
    tmp_dir = Path(tmp_dir)

    path_inverse = tmp_dir / "inverse_identification_without_raw_data.mat"

    robot_mat2hdf(save_path,path_inverse)
    

# %% ../../nbs/datasets/industrial_robot.ipynb 11
u_inverse = [ f'u{i}' for i in range(0,12) ]
y_inverse = [ f'y{i}' for i in range(0,6) ]

BenchmarkRobotInverse_Simulation = idb.BenchmarkSpecSimulation(
    name='BenchmarkRobotInverse_Simulation', dataset_id='robot_inverse',
    u_cols=u_inverse, y_cols=y_inverse, metric_func=identibench.metrics.rmse, 
    download_func=dl_robot_inverse,
    init_window=100
)
BenchmarkRobotInverse_Prediction = idb.BenchmarkSpecPrediction(
    name='BenchmarkRobotInverse_Prediction', dataset_id='robot_inverse',
    u_cols=u_inverse, y_cols=y_inverse, metric_func=identibench.metrics.rmse, 
    download_func=dl_robot_inverse,
    init_window=100, pred_horizon=100,pred_step=100
)
