from spinesTS.base import KerasMixin
import tensorflow as tf
from spinesTS.nn._keras_lambda import AddAxis
from tcn import TCN
from sklearn.base import RegressorMixin
from spinesTS.utils import seed_everything


class TCN1D(RegressorMixin, KerasMixin):
    def __init__(self, output_nums, learning_rate=0.1,
                 monitor='val_loss', min_delta=0, patience=10, restore_best_weights=True, random_seed=0):
        seed_everything(random_seed)
        self._output_nums = output_nums
        self.model = None
        self._learning_rate = learning_rate
        self._callback(monitor=monitor, min_delta=min_delta,
                       patience=patience, restore_best_weights=restore_best_weights)

    def _call(self):
        net = tf.keras.Sequential([
            tf.keras.Input(shape=self._shape),
            # tf.keras.layers.Dense(max(self._shape[0], self._output_nums), activation='relu', input_shape=self._shape),
            tf.keras.layers.BatchNormalization(),
            AddAxis(axis=-1),
            TCN(
                nb_filters=128,
                kernel_size=4,
                return_sequences=True
            ),
            tf.keras.layers.GlobalAveragePooling1D(),
            tf.keras.layers.Dense(16 * self._output_nums, activation='relu'),
            tf.keras.layers.Dense(self._output_nums)
        ])

        net.compile(metrics=['mae'], optimizer=tf.keras.optimizers.Adam(learning_rate=self._learning_rate)
                    , loss=tf.keras.losses.Huber()
                    )
        return net

    def fit(self, x, y, eval_set=None, epochs=1000, verbose=2, batch_size='auto'):
        self._fit_on_batch(x, batch_size=batch_size)
        self.model = self._call()
        self.model.fit(x, y, validation_data=eval_set, epochs=epochs,
                       verbose=verbose, batch_size=self._batch_size, callbacks=self._callback
                       )

    def predict(self, x):
        assert self.model is not None, "model not fitted yet."
        return self.model.predict(x)



tensorflow.keras.Input tensorflow.keras.layers.Dense tensorflow.keras.layers.GlobalAveragePooling1D tensorflow.keras.losses.Huber tensorflow.keras.layers.BatchNormalization tensorflow.keras.optimizers.Adam