"""
CNNbest Architecture
====================
The convolutional neural network architecture from the ASCAD paper
(Benadjila et al., 2020).

Architecture:
  Input(700,1) -> 5 Conv blocks [Conv1D -> SELU -> BatchNorm -> AvgPool]
  -> Flatten -> 2 x Dense(4096, SELU) -> Dense(256, Softmax)

Conv filters: [64, 128, 256, 512, 512], kernel_size=11, pool_size=2
Optimizer: RMSprop(lr=1e-5)
"""

from typing import Dict, Any, List

import tensorflow as tf

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


class CNNBest(BaseModel):
    """
    CNNbest architecture for ASCAD side-channel attacks.

    This is the best-performing CNN architecture identified in the original
    ASCAD paper. It uses 5 convolutional blocks with increasing filter counts,
    SELU activation, batch normalization, and average pooling, followed by
    two large fully connected layers.
    """

    def __init__(
        self,
        input_length: int = WINDOW_SIZE,
        num_classes: int = NUM_CLASSES,
        conv_filters: List[int] = None,
        kernel_size: int = CNN_DEFAULTS["kernel_size"],
        pool_size: int = CNN_DEFAULTS["pool_size"],
        fc_units: int = CNN_DEFAULTS["fc_units"],
        num_fc_layers: int = CNN_DEFAULTS["num_fc_layers"],
    ) -> None:
        """
        Args:
            input_length: Number of time samples in the input trace.
            num_classes: Number of output classes.
            conv_filters: List of filter counts for each conv block.
            kernel_size: Kernel size for all conv layers.
            pool_size: Pool size for average pooling.
            fc_units: Number of units in each fully connected layer.
            num_fc_layers: Number of fully connected layers.
        """
        super().__init__(input_shape=(input_length, 1), num_classes=num_classes)
        self.input_length = input_length
        self.conv_filters = conv_filters or list(CNN_DEFAULTS["conv_filters"])
        self.kernel_size = kernel_size
        self.pool_size = pool_size
        self.fc_units = fc_units
        self.num_fc_layers = num_fc_layers

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

        # Convolutional blocks
        for i, filters in enumerate(self.conv_filters):
            x = tf.keras.layers.Conv1D(
                filters=filters,
                kernel_size=self.kernel_size,
                padding="same",
                kernel_initializer="glorot_uniform",
                name=f"conv_{i}",
            )(x)
            x = tf.keras.layers.Activation("selu", name=f"selu_{i}")(x)
            x = tf.keras.layers.BatchNormalization(name=f"bn_{i}")(x)
            x = tf.keras.layers.AveragePooling1D(
                pool_size=self.pool_size, name=f"avgpool_{i}"
            )(x)

        x = tf.keras.layers.Flatten(name="flatten")(x)

        # Fully connected layers
        for i in range(self.num_fc_layers):
            x = tf.keras.layers.Dense(
                self.fc_units,
                activation="selu",
                kernel_initializer="glorot_uniform",
                name=f"fc_{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="CNNbest")
        return self._model

    def get_config(self) -> Dict[str, Any]:
        """Return architecture hyperparameters."""
        return {
            "model_type": "cnn",
            "architecture": "CNNbest",
            "input_length": self.input_length,
            "num_classes": self.num_classes,
            "conv_filters": self.conv_filters,
            "kernel_size": self.kernel_size,
            "pool_size": self.pool_size,
            "fc_units": self.fc_units,
            "num_fc_layers": self.num_fc_layers,
            "total_params": self.model.count_params(),
        }