"""FeedForwardNetwork definition required for deserializing model.keras.

This must be imported before tf.keras.models.load_model() is called so
that keras can resolve the registered custom class.
"""

import keras
import tensorflow as tf


@keras.saving.register_keras_serializable()
class FeedForwardNetwork(tf.keras.Model):
    """Fully-connected feedforward network for 6-class HAR classification.

    Architecture: Dense(512) → BN → ReLU → Dropout
                  Dense(256) → BN → ReLU → Dropout
                  Dense(128) → BN → ReLU → Dropout
                  Dense(6, softmax)
    """

    def __init__(
        self,
        num_features,
        num_classes,
        hidden_units=(512, 256, 128),
        dropout_rate=0.3,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self._num_features = num_features
        self._num_classes = num_classes
        self._hidden_units = tuple(hidden_units)
        self._dropout_rate = dropout_rate

        self.hidden_blocks = []
        for units in hidden_units:
            self.hidden_blocks.append([
                tf.keras.layers.Dense(units, use_bias=False),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.ReLU(),
                tf.keras.layers.Dropout(dropout_rate),
            ])

        self.output_layer = tf.keras.layers.Dense(num_classes, activation="softmax")

    def call(self, inputs, training=False):
        x = inputs
        for block in self.hidden_blocks:
            for layer in block:
                if isinstance(layer, (tf.keras.layers.BatchNormalization,
                                      tf.keras.layers.Dropout)):
                    x = layer(x, training=training)
                else:
                    x = layer(x)
        return self.output_layer(x)

    def get_config(self):
        config = super().get_config()
        config.update({
            "num_features": self._num_features,
            "num_classes": self._num_classes,
            "hidden_units": self._hidden_units,
            "dropout_rate": self._dropout_rate,
        })
        return config