"""Model classes required for deserializing .keras files.

Must be imported before tf.keras.models.load_model() so that keras
can resolve the registered custom classes.
"""

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


@keras.saving.register_keras_serializable()
class Conv1DNetwork(tf.keras.Model):
    """1D-CNN for classification on pre-computed feature vectors.

    Architecture: Reshape(561, 1)
                  → Conv1D(64,  k=5) → BN → ReLU → MaxPool(2) → Dropout(0.3)
                  → Conv1D(128, k=5) → BN → ReLU → MaxPool(2) → Dropout(0.3)
                  → Conv1D(256, k=3) → BN → ReLU → GlobalAvgPool1D
                  → Dense(128) → BN → ReLU → Dropout(0.5)
                  → Dense(6, softmax)
    """

    def __init__(
        self,
        num_features,
        num_classes,
        dropout_rate=0.3,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self._num_features = num_features
        self._num_classes = num_classes
        self._dropout_rate = dropout_rate

        self.reshape = tf.keras.layers.Reshape((num_features, 1))

        self.conv1 = tf.keras.layers.Conv1D(64, kernel_size=5, padding="same", use_bias=False)
        self.bn1   = tf.keras.layers.BatchNormalization()
        self.relu1 = tf.keras.layers.ReLU()
        self.pool1 = tf.keras.layers.MaxPooling1D(pool_size=2)
        self.drop1 = tf.keras.layers.Dropout(dropout_rate)

        self.conv2 = tf.keras.layers.Conv1D(128, kernel_size=5, padding="same", use_bias=False)
        self.bn2   = tf.keras.layers.BatchNormalization()
        self.relu2 = tf.keras.layers.ReLU()
        self.pool2 = tf.keras.layers.MaxPooling1D(pool_size=2)
        self.drop2 = tf.keras.layers.Dropout(dropout_rate)

        self.conv3 = tf.keras.layers.Conv1D(256, kernel_size=3, padding="same", use_bias=False)
        self.bn3   = tf.keras.layers.BatchNormalization()
        self.relu3 = tf.keras.layers.ReLU()
        self.gap   = tf.keras.layers.GlobalAveragePooling1D()

        self.dense1  = tf.keras.layers.Dense(128, use_bias=False)
        self.bn_fc   = tf.keras.layers.BatchNormalization()
        self.relu_fc = tf.keras.layers.ReLU()
        self.drop_fc = tf.keras.layers.Dropout(0.5)

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

    def call(self, inputs, training=False):
        x = self.reshape(inputs)

        x = self.conv1(x)
        x = self.bn1(x, training=training)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.drop1(x, training=training)

        x = self.conv2(x)
        x = self.bn2(x, training=training)
        x = self.relu2(x)
        x = self.pool2(x)
        x = self.drop2(x, training=training)

        x = self.conv3(x)
        x = self.bn3(x, training=training)
        x = self.relu3(x)
        x = self.gap(x)

        x = self.dense1(x)
        x = self.bn_fc(x, training=training)
        x = self.relu_fc(x)
        x = self.drop_fc(x, training=training)

        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,
            "dropout_rate": self._dropout_rate,
        })
        return config