# ATTRIBUTION: this file consolidates model descriptions originating from
# https://github.com/hfawaz/dl-4-tsc
# It was published on github.com by Ismail Fawaz, Hassan and Forestier, Germain
# and Weber, Jonathan and Idoumghar, Lhassane and Muller, Pierre-Alain
# for their work "Deep learning for time series classification: a review",
# in Data Mining and Knowledge Discovery.
# LICENSE is GNU GENERAL PUBLIC LICENSE version 3 (29 June 2007).


"""
Centralized model descriptions for Time Series Classification (TSC) models (v2).

This module collates callable model descriptions for the eleven TSC models:
- CNN (valid), CNN (same), FCN, MLP, RESNET, ENCODER
- INCEPTION, TLENET, MCDCNN, MCNN, TWIESN

Notes:
- Keras-based descriptions return a tuple: (compiled_keras_model, [callbacks]).
- Some legacy models (e.g., TWIESN) are not Keras graphs. For these, we expose
  a descriptor callable that raises NotImplementedError for Keras usage and
  provide a reference to the legacy class in the docstring for integration.
"""

import gc
import os
import pickle
import time
import numpy as np
from scipy import sparse
from scipy.sparse import linalg as slinalg
from sklearn.linear_model import Ridge
from typing import Callable, List, Tuple

import tensorflow.keras as keras



def _dl4tsc_cnn_constructor_helper(padding: str, input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    conv1 = keras.layers.Conv1D(filters=6, kernel_size=7, padding=padding, activation='sigmoid')(input_layer)
    conv1 = keras.layers.AveragePooling1D(pool_size=3)(conv1)

    conv2 = keras.layers.Conv1D(filters=12, kernel_size=7, padding=padding, activation='sigmoid')(conv1)
    conv2 = keras.layers.AveragePooling1D(pool_size=3)(conv2)

    flatten_layer = keras.layers.Flatten()(conv2)

    output_layer = keras.layers.Dense(units=nb_classes, activation='sigmoid')(flatten_layer)

    model = keras.models.Model(inputs=input_layer, outputs=output_layer)

    model.compile(loss='mean_squared_error', optimizer=keras.optimizers.Adam(), metrics=['accuracy'])

    return model, []


def dl4tsc_cnn_padding_valid(input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    return _dl4tsc_cnn_constructor_helper('valid', input_layer, nb_classes)


def dl4tsc_cnn_padding_same(input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    return _dl4tsc_cnn_constructor_helper('same', input_layer, nb_classes)


def dl4tsc_fcn(input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    conv1 = keras.layers.Conv1D(filters=128, kernel_size=8, padding='same')(input_layer)
    conv1 = keras.layers.BatchNormalization()(conv1)
    conv1 = keras.layers.Activation(activation='relu')(conv1)

    conv2 = keras.layers.Conv1D(filters=256, kernel_size=5, padding='same')(conv1)
    conv2 = keras.layers.BatchNormalization()(conv2)
    conv2 = keras.layers.Activation('relu')(conv2)

    conv3 = keras.layers.Conv1D(128, kernel_size=3, padding='same')(conv2)
    conv3 = keras.layers.BatchNormalization()(conv3)
    conv3 = keras.layers.Activation('relu')(conv3)

    gap_layer = keras.layers.GlobalAveragePooling1D()(conv3)

    output_layer = keras.layers.Dense(nb_classes, activation='softmax')(gap_layer)

    model = keras.models.Model(inputs=input_layer, outputs=output_layer)

    model.compile(loss='categorical_crossentropy', optimizer=keras.optimizers.Adam(), metrics=['accuracy'])

    reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor='loss', factor=0.5, patience=50, min_lr=0.0001)
    return model, [reduce_lr]


def dl4tsc_mlp(input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    input_layer_flattened = keras.layers.Flatten()(input_layer)

    layer_1 = keras.layers.Dropout(0.1)(input_layer_flattened)
    layer_1 = keras.layers.Dense(500, activation='relu')(layer_1)

    layer_2 = keras.layers.Dropout(0.2)(layer_1)
    layer_2 = keras.layers.Dense(500, activation='relu')(layer_2)

    layer_3 = keras.layers.Dropout(0.2)(layer_2)
    layer_3 = keras.layers.Dense(500, activation='relu')(layer_3)

    output_layer = keras.layers.Dropout(0.3)(layer_3)
    output_layer = keras.layers.Dense(nb_classes, activation='softmax')(output_layer)

    model = keras.models.Model(inputs=input_layer, outputs=output_layer)

    model.compile(loss='categorical_crossentropy', optimizer=keras.optimizers.Adadelta(), metrics=['accuracy'])

    reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor='loss', factor=0.5, patience=200, min_lr=0.1)

    return model, [reduce_lr]


def dl4tsc_resnet(input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    n_feature_maps = 64

    conv_x = keras.layers.Conv1D(filters=n_feature_maps, kernel_size=8, padding='same')(input_layer)
    conv_x = keras.layers.BatchNormalization()(conv_x)
    conv_x = keras.layers.Activation('relu')(conv_x)

    conv_y = keras.layers.Conv1D(filters=n_feature_maps, kernel_size=5, padding='same')(conv_x)
    conv_y = keras.layers.BatchNormalization()(conv_y)
    conv_y = keras.layers.Activation('relu')(conv_y)

    conv_z = keras.layers.Conv1D(filters=n_feature_maps, kernel_size=3, padding='same')(conv_y)
    conv_z = keras.layers.BatchNormalization()(conv_z)

    shortcut_y = keras.layers.Conv1D(filters=n_feature_maps, kernel_size=1, padding='same')(input_layer)
    shortcut_y = keras.layers.BatchNormalization()(shortcut_y)

    output_block_1 = keras.layers.add([shortcut_y, conv_z])
    output_block_1 = keras.layers.Activation('relu')(output_block_1)

    conv_x = keras.layers.Conv1D(filters=n_feature_maps * 2, kernel_size=8, padding='same')(output_block_1)
    conv_x = keras.layers.BatchNormalization()(conv_x)
    conv_x = keras.layers.Activation('relu')(conv_x)

    conv_y = keras.layers.Conv1D(filters=n_feature_maps * 2, kernel_size=5, padding='same')(conv_x)
    conv_y = keras.layers.BatchNormalization()(conv_y)
    conv_y = keras.layers.Activation('relu')(conv_y)

    conv_z = keras.layers.Conv1D(filters=n_feature_maps * 2, kernel_size=3, padding='same')(conv_y)
    conv_z = keras.layers.BatchNormalization()(conv_z)

    shortcut_y = keras.layers.Conv1D(filters=n_feature_maps * 2, kernel_size=1, padding='same')(output_block_1)
    shortcut_y = keras.layers.BatchNormalization()(shortcut_y)

    output_block_2 = keras.layers.add([shortcut_y, conv_z])
    output_block_2 = keras.layers.Activation('relu')(output_block_2)

    conv_x = keras.layers.Conv1D(filters=n_feature_maps * 2, kernel_size=8, padding='same')(output_block_2)
    conv_x = keras.layers.BatchNormalization()(conv_x)
    conv_x = keras.layers.Activation('relu')(conv_x)

    conv_y = keras.layers.Conv1D(filters=n_feature_maps * 2, kernel_size=5, padding='same')(conv_x)
    conv_y = keras.layers.BatchNormalization()(conv_y)
    conv_y = keras.layers.Activation('relu')(conv_y)

    conv_z = keras.layers.Conv1D(filters=n_feature_maps * 2, kernel_size=3, padding='same')(conv_y)
    conv_z = keras.layers.BatchNormalization()(conv_z)

    shortcut_y = keras.layers.BatchNormalization()(output_block_2)

    output_block_3 = keras.layers.add([shortcut_y, conv_z])
    output_block_3 = keras.layers.Activation('relu')(output_block_3)

    gap_layer = keras.layers.GlobalAveragePooling1D()(output_block_3)

    output_layer = keras.layers.Dense(nb_classes, activation='softmax')(gap_layer)

    model = keras.models.Model(inputs=input_layer, outputs=output_layer)

    model.compile(loss='categorical_crossentropy', optimizer=keras.optimizers.Adam(), metrics=['accuracy'])

    reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor='loss', factor=0.5, patience=50, min_lr=0.0001)

    return model, [reduce_lr]


def dl4tsc_encoder(input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    from useckit.paradigms.time_series_classification.prediction_models.instance_normalization_layer import InstanceNormalization
    conv1 = keras.layers.Conv1D(filters=128, kernel_size=5, strides=1, padding='same')(input_layer)
    conv1 = InstanceNormalization()(conv1)
    conv1 = keras.layers.PReLU(shared_axes=[1])(conv1)
    conv1 = keras.layers.Dropout(rate=0.2)(conv1)
    conv1 = keras.layers.MaxPooling1D(pool_size=2)(conv1)
    conv2 = keras.layers.Conv1D(filters=256, kernel_size=11, strides=1, padding='same')(conv1)
    conv2 = InstanceNormalization()(conv2)
    conv2 = keras.layers.PReLU(shared_axes=[1])(conv2)
    conv2 = keras.layers.Dropout(rate=0.2)(conv2)
    conv2 = keras.layers.MaxPooling1D(pool_size=2)(conv2)
    conv3 = keras.layers.Conv1D(filters=512, kernel_size=21, strides=1, padding='same')(conv2)
    conv3 = InstanceNormalization()(conv3)
    conv3 = keras.layers.PReLU(shared_axes=[1])(conv3)
    conv3 = keras.layers.Dropout(rate=0.2)(conv3)
    attention_data = keras.layers.Lambda(lambda x: x[:, :, :256])(conv3)
    attention_softmax = keras.layers.Lambda(lambda x: x[:, :, 256:])(conv3)
    attention_softmax = keras.layers.Softmax()(attention_softmax)
    multiply_layer = keras.layers.Multiply()([attention_softmax, attention_data])
    dense_layer = keras.layers.Dense(units=256, activation='sigmoid')(multiply_layer)
    dense_layer = InstanceNormalization()(dense_layer)
    flatten_layer = keras.layers.Flatten()(dense_layer)
    output_layer = keras.layers.Dense(units=nb_classes, activation='softmax')(flatten_layer)

    model = keras.models.Model(inputs=input_layer, outputs=output_layer)

    model.compile(loss='categorical_crossentropy', optimizer=keras.optimizers.Adam(0.00001), metrics=['accuracy'])

    return model, []


# Additional TSC models implemented here for v2
def dl4tsc_inception(input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    nb_filters = 32
    use_residual = True
    use_bottleneck = True
    depth = 6
    kernel_size = 41 - 1  # follow legacy impl offset
    bottleneck_size = 32

    def inception_module(x, stride=1, activation='linear'):
        if use_bottleneck and int(x.shape[-1]) > bottleneck_size:
            input_inception = keras.layers.Conv1D(
                filters=bottleneck_size, kernel_size=1, padding='same', activation=activation, use_bias=False
            )(x)
        else:
            input_inception = x

        kernel_size_s = [kernel_size // (2 ** i) for i in range(3)]
        conv_list = [
            keras.layers.Conv1D(filters=nb_filters, kernel_size=ks, strides=stride,
                                 padding='same', activation=activation, use_bias=False)(input_inception)
            for ks in kernel_size_s
        ]

        max_pool_1 = keras.layers.MaxPool1D(pool_size=3, strides=stride, padding='same')(x)
        conv_6 = keras.layers.Conv1D(filters=nb_filters, kernel_size=1,
                                     padding='same', activation=activation, use_bias=False)(max_pool_1)
        conv_list.append(conv_6)

        x_out = keras.layers.Concatenate(axis=2)(conv_list)
        x_out = keras.layers.BatchNormalization()(x_out)
        x_out = keras.layers.Activation('relu')(x_out)
        return x_out

    def shortcut_layer(x_in, x_out):
        shortcut_y = keras.layers.Conv1D(filters=int(x_out.shape[-1]), kernel_size=1,
                                         padding='same', use_bias=False)(x_in)
        shortcut_y = keras.layers.BatchNormalization()(shortcut_y)
        x = keras.layers.Add()([shortcut_y, x_out])
        x = keras.layers.Activation('relu')(x)
        return x

    x = input_layer
    input_res = input_layer
    for d in range(depth):
        x = inception_module(x)
        if use_residual and d % 3 == 2:
            x = shortcut_layer(input_res, x)
            input_res = x

    gap_layer = keras.layers.GlobalAveragePooling1D()(x)
    output_layer = keras.layers.Dense(nb_classes, activation='softmax')(gap_layer)
    model = keras.models.Model(inputs=input_layer, outputs=output_layer)

    model.compile(loss='categorical_crossentropy', optimizer=keras.optimizers.Adam(0.001), metrics=['accuracy'])
    reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor='loss', factor=0.5, patience=50, min_lr=0.0001)
    return model, [reduce_lr]


def dl4tsc_tlenet(input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    conv_1 = keras.layers.Conv1D(filters=5, kernel_size=5, activation='relu', padding='same')(input_layer)
    conv_1 = keras.layers.MaxPool1D(pool_size=2)(conv_1)

    conv_2 = keras.layers.Conv1D(filters=20, kernel_size=5, activation='relu', padding='same')(conv_1)
    conv_2 = keras.layers.MaxPool1D(pool_size=4)(conv_2)

    flatten_layer = keras.layers.Flatten()(conv_2)
    fully_connected_layer = keras.layers.Dense(500, activation='relu')(flatten_layer)

    output_layer = keras.layers.Dense(nb_classes, activation='softmax')(fully_connected_layer)
    model = keras.models.Model(inputs=input_layer, outputs=output_layer)

    model.compile(optimizer=keras.optimizers.Adam(learning_rate=0.01, decay=0.005),
                  loss='categorical_crossentropy', metrics=['accuracy'])
    return model, []


def dl4tsc_mcdcnn(input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    n_t = int(input_layer.shape[1])
    n_vars = int(input_layer.shape[2]) if len(input_layer.shape) > 2 else 1

    padding = 'valid'
    if n_t < 60:
        padding = 'same'

    input_layers = []
    conv2_layers = []
    for _ in range(n_vars):
        in_i = keras.layers.Input((n_t, 1))
        input_layers.append(in_i)
        c1 = keras.layers.Conv1D(filters=8, kernel_size=5, activation='relu', padding=padding)(in_i)
        c1 = keras.layers.MaxPooling1D(pool_size=2)(c1)
        c2 = keras.layers.Conv1D(filters=8, kernel_size=5, activation='relu', padding=padding)(c1)
        c2 = keras.layers.MaxPooling1D(pool_size=2)(c2)
        c2 = keras.layers.Flatten()(c2)
        conv2_layers.append(c2)

    if n_vars == 1:
        concat_layer = conv2_layers[0]
    else:
        concat_layer = keras.layers.Concatenate(axis=-1)(conv2_layers)

    fully_connected = keras.layers.Dense(units=732, activation='relu')(concat_layer)
    output_layer = keras.layers.Dense(nb_classes, activation='softmax')(fully_connected)
    model = keras.models.Model(inputs=input_layers, outputs=output_layer)

    # Match legacy optimizer config
    model.compile(loss='categorical_crossentropy',
                  optimizer=keras.optimizers.SGD(learning_rate=0.01, momentum=0.9, decay=0.0005),
                  metrics=['accuracy'])
    return model, []


def dl4tsc_mcnn(input_layer, nb_classes) -> Tuple["keras.Model", List[Callable]]:  
    # Derive number of variables from input layer shape (T, V)
    n_t = int(input_layer.shape[1])
    n_vars = int(input_layer.shape[2]) if len(input_layer.shape) > 2 else 1

    pool_factor = 2
    kernel_size = max(3, min(11, n_t // 10))

    stage_1_layers = []
    input_layers = []
    for _ in range(n_vars):
        in_i = keras.layers.Input((n_t, 1))
        input_layers.append(in_i)
        conv_layer = keras.layers.Conv1D(filters=256, kernel_size=kernel_size, padding='same', activation='sigmoid')(in_i)
        pool_size = max(1, int(int(conv_layer.shape[1]) / pool_factor))
        max_layer = keras.layers.MaxPooling1D(pool_size=pool_size)(conv_layer)
        stage_1_layers.append(max_layer)

    concat_layer = stage_1_layers[0] if n_vars == 1 else keras.layers.Concatenate(axis=-1)(stage_1_layers)
    kernel_size_2 = int(min(kernel_size, int(concat_layer.shape[1])))
    full_conv = keras.layers.Conv1D(filters=256, kernel_size=kernel_size_2, padding='same', activation='sigmoid')(concat_layer)
    pool_size_2 = max(1, int(int(full_conv.shape[1]) / pool_factor))
    full_max = keras.layers.MaxPooling1D(pool_size=pool_size_2)(full_conv)
    full_max = keras.layers.Flatten()(full_max)
    fully_connected = keras.layers.Dense(units=256, activation='sigmoid')(full_max)

    output_layer = keras.layers.Dense(nb_classes, activation='softmax')(fully_connected)
    model = keras.models.Model(inputs=input_layers, outputs=output_layer)
    # Follow model_builder’s config for MCNN
    model.compile(loss='categorical_crossentropy', optimizer=keras.optimizers.Adam(0.1), metrics=['accuracy'])
    return model, []


def _legacy_create_directory(path: str):
    """Create a directory if it does not exist (legacy helper).

    Mirrors the minimal behavior of the legacy `create_directory` util used in
    the original TSC codebase. Safe to call repeatedly.
    """
    os.makedirs(path, exist_ok=True)


class dl4tsc_twiesn:
    def __init__(self, output_dir: str = "TWIESN_model_out"):
        # This class is a lightweight, self-contained port of the legacy TWIESN
        # model for v2. It intentionally avoids inheriting from legacy bases.
        self.output_dir = output_dir
        # hyperparameters
        first_config = {'N_x': 250, 'connect': 0.5, 'scaleW_in': 1.0, 'lamda': 0.0}
        second_config = {'N_x': 250, 'connect': 0.5, 'scaleW_in': 2.0, 'lamda': 0.05}
        third_config = {'N_x': 500, 'connect': 0.1, 'scaleW_in': 2.0, 'lamda': 0.05}
        fourth_config = {'N_x': 800, 'connect': 0.1, 'scaleW_in': 2.0, 'lamda': 0.05}
        self.configs = [first_config, second_config, third_config, fourth_config]
        self.rho_s = [0.55, 0.9, 2.0, 5.0]
        self.alpha = 0.1  # leaky rate
        self.y_pred = None
        self.y_pred_reshaped = None

    def persist_model_trainables(self, path_like: str = 'saved_trainable_values', **kwargs):
        raise NotImplementedError()

    def restore_model_trainables(self, path_like, dataset = None, **kwargs):
        raise NotImplementedError()

    def init_matrices(self):
        self.W_in = (2.0 * np.random.rand(self.N_x, self.num_dim) - 1.0) / (2.0 * self.scaleW_in)

        converged = False
        i = 0
        while not converged:
            i += 1
            self.W = sparse.rand(self.N_x, self.N_x, density=self.connect).todense()
            self.W[np.where(self.W > 0)] -= 0.5
            try:
                eig, _ = slinalg.eigs(self.W, k=1, which='LM')
                converged = True
            except Exception:
                print('not converged ', i)
                continue
        self.W /= np.abs(eig) / self.rho

    def compute_state_matrix(self, x_in):
        n = x_in.shape[0]
        X_t = np.zeros((n, self.T, self.N_x), dtype=np.float64)
        X_t_1 = np.zeros((n, self.N_x), dtype=np.float64)
        for t in range(self.T):
            curr_in = x_in[:, t, :]
            curr_state = np.tanh(self.W_in.dot(curr_in.T) + self.W.dot(X_t_1.T)).T
            curr_state = (1 - self.alpha) * X_t_1 + self.alpha * curr_state
            X_t_1 = curr_state
            X_t[:, t, :] = curr_state
        return X_t

    def reshape_prediction(self, y_pred, num_instances, length_series):
        new_y_pred = y_pred.reshape(num_instances, length_series, y_pred.shape[-1])
        new_y_pred = np.average(new_y_pred, axis=1)
        new_y_pred = np.argmax(new_y_pred, axis=1)
        return new_y_pred

    def train(self):
        start_time = time.time()
        self.init_matrices()
        state_matrix = self.compute_state_matrix(self.x_train)
        new_x_train = np.concatenate((self.x_train, state_matrix), axis=2).reshape(
            self.N * self.T, self.num_dim + self.N_x)
        state_matrix = None
        gc.collect()
        new_labels = np.repeat(self.y_train, self.T, axis=0)
        ridge_classifier = Ridge(alpha=self.lamda)
        ridge_classifier.fit(new_x_train, new_labels)

        state_matrix = self.compute_state_matrix(self.x_val)
        new_x_val = np.concatenate((self.x_val, state_matrix), axis=2).reshape(
            self.x_val.shape[0] * self.T, self.num_dim + self.N_x)
        y_pred_val = ridge_classifier.predict(new_x_val)
        y_pred_val = self.reshape_prediction(y_pred_val, self.x_val.shape[0], self.T)
        df_val_metrics = self.calculate_metrics(np.argmax(self.y_val, axis=1), y_pred_val, 0.0)
        train_acc = df_val_metrics['accuracy'][0]

        state_matrix = self.compute_state_matrix(self.x_test)
        new_x_test = np.concatenate((self.x_test, state_matrix), axis=2).reshape(self.x_test.shape[0] * self.T,
                                                                                 self.num_dim + self.N_x)
        state_matrix = None
        gc.collect()
        self.y_pred = y_pred = ridge_classifier.predict(new_x_test)
        self.y_pred_reshaped = y_pred_reshaped = self.reshape_prediction(y_pred, self.x_test.shape[0], self.T)

        duration = time.time() - start_time
        df_metrics = self.calculate_metrics(self.y_true, y_pred_reshaped, duration)
        self.W_out = ridge_classifier.coef_
        gc.collect()
        np.savetxt(os.path.join(self.output_dir, 'W_in.txt'), self.W_in)
        np.savetxt(os.path.join(self.output_dir, 'W.txt'), self.W)
        np.savetxt(os.path.join(self.output_dir, 'W_out.txt'), self.W_out)
        df_metrics.to_csv(os.path.join(self.output_dir, 'df_metrics.csv'), index=False)
        return df_metrics, train_acc, ridge_classifier

    def convert_dataset_to_legacy_values(self, dataset) -> (
        np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int
    ):
        """Convert v2 Dataset into legacy tuple expected by dl4tsc implementations.

        Returns (x_train, x_val, y_train, y_val, y_true, input_shape, nb_classes).
        """
        # Import locally to avoid hard dependency at module import time
        from useckit.v2.dataset import Dataset as V2Dataset

        if not isinstance(dataset, V2Dataset):
            # Be permissive: assume a duck-typed object with the same interface
            pass

        x_train, x_val = dataset.trainset_data, dataset.validationset_data
        y_train, y_val, _, _ = dataset.view_one_hot_encoded_labels()

        y_true = np.argmax(y_val, axis=1)
        nb_classes = y_train.shape[1] if y_train is not None else int(np.max(y_true) + 1)
        input_shape = x_train.shape[1:]

        assert isinstance(nb_classes, int)
        assert x_train.shape[0] == y_train.shape[0]
        assert x_val.shape[0] == y_val.shape[0]

        return x_train, x_val, y_train, y_val, y_true, input_shape, nb_classes

    def fit(self, dataset, *args, **kwargs):
        self.x_train, self.x_val, self.y_train, self.y_val, y_true, input_shape, nb_classes = self.convert_dataset_to_legacy_values(dataset)

        best_train_acc = -1

        self.num_dim = self.x_train.shape[2]
        self.T = self.x_train.shape[1]
        self.x_test = self.x_val
        self.y_true = y_true
        self.y_test = self.y_val

        self.N = self.x_train.shape[0]

        if self.x_train.shape[0] > 1000 or self.x_test.shape[0] > 1000:
            for config in self.configs:
                config['N_x'] = 100
            self.configs = [self.configs[0], self.configs[1], self.configs[2]]

        output_directory_root = self.output_dir
        for idx_config in range(len(self.configs)):
            for rho in self.rho_s:
                self.rho = rho
                self.N_x = self.configs[idx_config]['N_x']
                self.connect = self.configs[idx_config]['connect']
                self.scaleW_in = self.configs[idx_config]['scaleW_in']
                self.lamda = self.configs[idx_config]['lamda']
                self.output_dir = os.path.join(output_directory_root,
                                               'hyper_param_search',
                                               'config_' + str(idx_config),
                                               'rho_' + str(rho))
                _legacy_create_directory(self.output_dir)
                df_metrics, train_acc, clf = self.train()

                if best_train_acc < train_acc:
                    best_train_acc = train_acc
                    self.y_pred_best = self.y_pred
                    self.y_pred_reshaped_best = self.y_pred_reshaped
                    df_metrics.to_csv(os.path.join(output_directory_root, 'df_metrics.csv'), index=False)
                    np.savetxt(os.path.join(output_directory_root, 'W_in.txt'), self.W_in)
                    np.savetxt(os.path.join(output_directory_root, 'W.txt'), self.W)
                    np.savetxt(os.path.join(output_directory_root, 'W_out.txt'), self.W_out)
                    with open(os.path.join(output_directory_root, 'ridge.pkl'), 'wb') as file:
                        pickle.dump(clf, file)
                    with open(os.path.join(output_directory_root, 'T.pkl'), 'wb') as file:
                        pickle.dump(self.T, file)
                    with open(os.path.join(output_directory_root, 'N_x.pkl'), 'wb') as file:
                        pickle.dump(self.N_x, file)
                    with open(os.path.join(output_directory_root, 'alpha.pkl'), 'wb') as file:
                        pickle.dump(self.alpha, file)

                gc.collect()

    def predict(self, x_test, y_true, x_train, y_train, y_test, return_df_metrics=True):
        from pathlib import Path
        modelroot = Path(self.output_dir).parent.parent.parent
        with open(modelroot.joinpath("ridge.pkl"), 'rb') as file:
            ridge_classifier = pickle.load(file)
        with open(modelroot.joinpath("T.pkl"), 'rb') as file:
            self.T = pickle.load(file)
        with open(modelroot.joinpath("N_x.pkl"), 'rb') as file:
            self.N_x = pickle.load(file)
        with open(modelroot.joinpath("alpha.pkl"), 'rb') as file:
            self.alpha = pickle.load(file)
        self.W_in = np.loadtxt(modelroot.joinpath('W_in.txt'))
        self.W = np.loadtxt(modelroot.joinpath('W.txt'))
        self.W_out = np.loadtxt(modelroot.joinpath('W_out.txt'))

        state_matrix = self.compute_state_matrix(x_test)
        new_x_test = np.concatenate((x_test, state_matrix), axis=2).reshape(x_test.shape[0] * self.T,
                                                                            self.num_dim + self.N_x)
        state_matrix = None
        gc.collect()
        y_pred = ridge_classifier.predict(new_x_test)
        y_pred_reshaped = self.reshape_prediction(y_pred, x_test.shape[0], self.T)

        return y_pred_reshaped

    def calculate_metrics(self, y_true, y_pred, duration, y_true_val=None, y_pred_val=None):
        import pandas as pd
        from sklearn.metrics import precision_score, accuracy_score, recall_score

        res = pd.DataFrame(data=np.zeros((1, 4), dtype=np.float64), index=[0],
                           columns=['precision', 'accuracy', 'recall', 'duration'])

        res['precision'] = precision_score(y_true, y_pred, average='macro')
        res['accuracy'] = accuracy_score(y_true, y_pred)

        if not y_true_val is None:
            # this is useful when transfer learning is used with cross validation
            res['accuracy_val'] = accuracy_score(y_true_val, y_pred_val)

        res['recall'] = recall_score(y_true, y_pred, average='macro')
        res['duration'] = duration
        return res



MODEL_DESCRIPTIONS = {
    'CNN_VALID': dl4tsc_cnn_padding_valid,
    'CNN_SAME': dl4tsc_cnn_padding_same,
    'FCN': dl4tsc_fcn,
    'MLP': dl4tsc_mlp,
    'RESNET': dl4tsc_resnet,
    'ENCODER': dl4tsc_encoder,
    'INCEPTION': dl4tsc_inception,
    'TLENET': dl4tsc_tlenet,
    'MCDCNN': dl4tsc_mcdcnn,
    'MCNN': dl4tsc_mcnn,
    # Non-Keras
    'TWIESN': dl4tsc_twiesn,
}
