"""
MLPbest Architecture
====================
The multi-layer perceptron architecture from the ASCAD paper (Benadjila et al., 2020).

Architecture: Input(700) -> 5 x Dense(200, ReLU) -> Dense(256, Softmax)
Optimizer: RMSprop(lr=1e-5)
"""

from typing import Dict, Any

import tensorflow as tf

from .base import BaseModel
from ..constants import MLP_DEFAULTS, NUM_CLASSES, WINDOW_SIZE


class MLPBest(BaseModel):
    """
    MLPbest architecture for ASCAD side-channel attacks.

    This is the best-performing MLP architecture identified in the original
    ASCAD paper. It consists of 5 hidden dense layers with 200 units each,
    followed by a softmax classification layer over 256 S-Box output classes.
    """

    def __init__(
        self,
        input_length: int = WINDOW_SIZE,
        num_classes: int = NUM_CLASSES,
        num_hidden_layers: int = MLP_DEFAULTS["num_hidden_layers"],
        hidden_units: int = MLP_DEFAULTS["hidden_units"],
    ) -> None:
        """
        Args:
            input_length: Number of time samples in the input trace.
            num_classes: Number of output classes.
            num_hidden_layers: Number of hidden dense layers.
            hidden_units: Number of units per hidden layer.
        """
        super().__init__(input_shape=(input_length,), num_classes=num_classes)
        self.input_length = input_length
        self.num_hidden_layers = num_hidden_layers
        self.hidden_units = hidden_units

    def build(self) -> tf.keras.Model:
        """Construct the MLPbest Keras model."""
        inputs = tf.keras.Input(shape=self.input_shape, name="trace_input")
        x = inputs

        for i in range(self.num_hidden_layers):
            x = tf.keras.layers.Dense(
                self.hidden_units,
                activation="relu",
                kernel_initializer="glorot_uniform",
                name=f"dense_{i}",
            )(x)

        outputs = tf.keras.layers.Dense(
            self.num_classes,
            activation="softmax",
            name="predictions",
        )(x)

        self._model = tf.keras.Model(inputs=inputs, outputs=outputs, name="MLPbest")
        return self._model

    def get_config(self) -> Dict[str, Any]:
        """Return architecture hyperparameters."""
        return {
            "model_type": "mlp",
            "architecture": "MLPbest",
            "input_length": self.input_length,
            "num_classes": self.num_classes,
            "num_hidden_layers": self.num_hidden_layers,
            "hidden_units": self.hidden_units,
            "total_params": self.model.count_params(),
        }