import tensorflow as tf
from tensorflow.keras.layers import Layer, Dense

# ---------------------------------------------------------------------------
# Module-level activation map (single source of truth)
# ---------------------------------------------------------------------------
ACTIVATION_MAP = {
    "relu": tf.nn.relu,
    "sigmoid": tf.nn.sigmoid,
    "tanh": tf.nn.tanh,
    "softmax": tf.nn.softmax,
    "softplus": tf.nn.softplus,
    "softsign": tf.nn.softsign,
    "elu": tf.nn.elu,
    "selu": tf.nn.selu,
    "swish": tf.nn.swish,
    "gelu": tf.nn.gelu,
    "leaky_relu": tf.nn.leaky_relu,
    "relu6": tf.nn.relu6,
    "hard_sigmoid": tf.keras.activations.hard_sigmoid,
    "exponential": tf.keras.activations.exponential,
    "linear": tf.keras.activations.linear,
    "log_softmax": tf.nn.log_softmax,
}


def _resolve_activations(activations):
    """Resolve a list of activation names/functions into (functions, names)."""
    funcs, names = [], []
    for act in activations:
        if isinstance(act, str):
            if act not in ACTIVATION_MAP:
                raise ValueError(f"Unknown activation function name: {act}")
            funcs.append(ACTIVATION_MAP[act])
            names.append(act)
        else:
            funcs.append(act)
            names.append(act.__name__)
    return funcs, names


# ---------------------------------------------------------------------------
# WAFLayer - Weighted Activation Fusion
# ---------------------------------------------------------------------------
@tf.keras.utils.register_keras_serializable()
class WAFLayer(Layer):
    def __init__(self, units=32, activations=None, compressor="sum",
                 apply_father_activation=False, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        self.compressor = compressor
        self.apply_father_activation = apply_father_activation

        activations = activations or ["relu", "swish"]
        self.activations, self.activation_names = _resolve_activations(activations)

    def build(self, input_shape):
        n_act = len(self.activations)
        self.w = self.add_weight(
            name="kernel",
            shape=(input_shape[-1], self.units),
            initializer="glorot_uniform",
            trainable=True,
        )
        self.wf = self.add_weight(
            name="kernel_functions",
            shape=(self.units, n_act),
            initializer="glorot_uniform",
            trainable=True,
        )
        self.b = self.add_weight(
            name="bias",
            shape=(self.units,),
            initializer="zeros",
            trainable=True,
        )
        self.bf = self.add_weight(
            name="bias_functions",
            shape=(self.units, n_act),
            initializer="zeros",
            trainable=True,
        )
        if self.apply_father_activation:
            self.softingparam = self.add_weight(
                name="softingparam",
                shape=(1,),
                initializer=tf.keras.initializers.Constant(1.0),
                trainable=True,
            )
        super().build(input_shape)

    def call(self, inputs):
        weighted_sum = tf.matmul(inputs, self.w) + self.b
        expanded_ws = tf.expand_dims(weighted_sum, axis=-1)
        expanded_ws = expanded_ws * self.wf + self.bf

        results = [func(expanded_ws[..., i]) for i, func in enumerate(self.activations)]

        result = self._compress(results)

        if self.apply_father_activation:
            result = result * tf.nn.sigmoid(self.softingparam * result)

        return result

    def _compress(self, results):
        stacked = tf.stack(results, axis=-1)
        if self.compressor == "sum":
            return tf.reduce_sum(stacked, axis=-1)
        elif self.compressor == "avg":
            return tf.reduce_mean(stacked, axis=-1)
        else:
            raise ValueError("Invalid compressor. Use 'sum' or 'avg'.")

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.units)

    def get_config(self):
        config = super().get_config()
        config.update({
            "units": self.units,
            "activations": self.activation_names,
            "compressor": self.compressor,
            "apply_father_activation": self.apply_father_activation,
        })
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


# ---------------------------------------------------------------------------
# WAFLayerNorm - WAF with pre-activation normalization (V1)
# ---------------------------------------------------------------------------
@tf.keras.utils.register_keras_serializable()
class WAFLayerNorm(WAFLayer):
    """WAFLayer that normalizes the linear projection before branching."""

    def call(self, inputs):
        weighted_sum = tf.matmul(inputs, self.w) + self.b

        # Normalize pre-activation
        mean, variance = tf.nn.moments(weighted_sum, axes=[-1], keepdims=True)
        normalized = (weighted_sum - mean) / tf.sqrt(variance + 1e-5)

        expanded_norm = tf.expand_dims(normalized, axis=-1)
        transformed = expanded_norm * self.wf + self.bf

        results = [func(transformed[..., i]) for i, func in enumerate(self.activations)]

        result = self._compress(results)

        if self.apply_father_activation:
            result = result * tf.nn.sigmoid(self.softingparam * result)

        return result


# ---------------------------------------------------------------------------
# WAFLayerNormV2 - WAF with post-activation normalization per branch
# ---------------------------------------------------------------------------
@tf.keras.utils.register_keras_serializable()
class WAFLayerNormV2(WAFLayer):
    """WAFLayer that normalizes each activation output individually before combining.

    This ensures activations with different output ranges (e.g., exponential vs tanh)
    are on the same scale before fusion.
    """

    def call(self, inputs):
        weighted_sum = tf.matmul(inputs, self.w) + self.b

        expanded_sum = tf.expand_dims(weighted_sum, axis=-1)
        transformed = expanded_sum * self.wf + self.bf

        results = [func(transformed[..., i]) for i, func in enumerate(self.activations)]

        # Normalize each activation output individually
        normalized_results = []
        for r in results:
            mean, variance = tf.nn.moments(r, axes=[-1], keepdims=True)
            r_norm = (r - mean) / tf.sqrt(variance + 1e-5)
            normalized_results.append(r_norm)

        result = self._compress(normalized_results)

        if self.apply_father_activation:
            result = result * tf.nn.sigmoid(self.softingparam * result)

        return result


# ---------------------------------------------------------------------------
# MAXLayerWithAttention - Mixture-of-Activation with input-dependent gating (USM)
# ---------------------------------------------------------------------------
@tf.keras.utils.register_keras_serializable()
class MAXLayerWithAttention(Layer):
    def __init__(self, units=32, activations=None, return_sequences=True, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        self.return_sequences = return_sequences

        activations = activations or ["relu", "swish"]
        self.activations, self.activation_names = _resolve_activations(activations)

        self.attention_dense = Dense(len(self.activations), activation="softmax")

    def build(self, input_shape):
        self.input_dim = input_shape[-1]
        n_act = len(self.activations)
        self.w = self.add_weight(
            name="kernel",
            shape=(self.input_dim, self.units),
            initializer="glorot_uniform",
            trainable=True,
        )
        self.b = self.add_weight(
            name="bias",
            shape=(self.units,),
            initializer="zeros",
            trainable=True,
        )
        self.wf = self.add_weight(
            name="kernel_functions",
            shape=(self.units, n_act),
            initializer="glorot_uniform",
            trainable=True,
        )
        self.bf = self.add_weight(
            name="bias_functions",
            shape=(self.units, n_act),
            initializer="zeros",
            trainable=True,
        )
        self.activation_counts = self.add_weight(
            name="activation_counts",
            shape=(n_act,),
            initializer="zeros",
            trainable=False,
        )
        self.attention_dense.build(input_shape)
        super().build(input_shape)

    def call(self, inputs):
        input_is_2d = (len(inputs.shape) == 2)
        if input_is_2d:
            inputs = tf.expand_dims(inputs, axis=1)

        weighted_sum = tf.matmul(inputs, self.w) + self.b
        expanded_ws = tf.expand_dims(weighted_sum, axis=-1)
        weighted_activations = expanded_ws * self.wf + self.bf

        results = [func(weighted_activations[..., i]) for i, func in enumerate(self.activations)]
        result_activations = tf.stack(results, axis=-1)

        attention_weights = self.attention_dense(inputs)
        attention_weights_expanded = tf.expand_dims(attention_weights, axis=2)
        weighted_results = result_activations * attention_weights_expanded
        final_results = tf.reduce_sum(weighted_results, axis=-1)

        if not self.return_sequences:
            final_results = tf.reduce_mean(final_results, axis=1)

        if input_is_2d and self.return_sequences:
            final_results = tf.squeeze(final_results, axis=1)

        activation_selection = tf.reduce_sum(attention_weights, axis=[0, 1])
        self.activation_counts.assign_add(activation_selection)

        return final_results

    def compute_output_shape(self, input_shape):
        if self.return_sequences:
            return (input_shape[0], input_shape[1], self.units)
        return (input_shape[0], self.units)

    def get_config(self):
        config = super().get_config()
        config.update({
            "units": self.units,
            "activations": self.activation_names,
            "return_sequences": self.return_sequences,
        })
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


# ---------------------------------------------------------------------------
# MAXLayerWithSelfAttention - Mixture-of-Activation with self-attention gating (GSM V1)
# ---------------------------------------------------------------------------
@tf.keras.utils.register_keras_serializable()
class MAXLayerWithSelfAttention(Layer):
    def __init__(self, units=32, activations=None, return_sequences=True, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        self.return_sequences = return_sequences

        activations = activations or ["relu", "swish"]
        self.activations, self.activation_names = _resolve_activations(activations)

        self.q_dense = Dense(units)
        self.k_dense = Dense(units)
        self.v_dense = Dense(units)
        self.activation_weight_dense = Dense(len(self.activations), activation="softmax")

    def build(self, input_shape):
        self.input_dim = input_shape[-1]
        n_act = len(self.activations)

        self.w = self.add_weight(
            name="kernel",
            shape=(self.input_dim, self.units),
            initializer="glorot_uniform",
            trainable=True,
        )
        self.b = self.add_weight(
            name="bias",
            shape=(self.units,),
            initializer="zeros",
            trainable=True,
        )
        self.activation_counts = self.add_weight(
            name="activation_counts",
            shape=(n_act,),
            initializer="zeros",
            trainable=False,
        )

        self.q_dense.build((None, None, self.units))
        self.k_dense.build((None, None, self.units))
        self.v_dense.build((None, None, self.units))
        self.activation_weight_dense.build((None, self.units))

        super().build(input_shape)

    def call(self, inputs):
        input_is_2d = (len(inputs.shape) == 2)
        if input_is_2d:
            inputs = tf.expand_dims(inputs, axis=1)

        x = tf.matmul(inputs, self.w) + self.b

        Q = self.q_dense(x)
        K = self.k_dense(x)
        V = self.v_dense(x)

        scores = tf.matmul(Q, K, transpose_b=True) / tf.sqrt(tf.cast(self.units, tf.float32))
        attention_weights = tf.nn.softmax(scores, axis=-1)
        context = tf.matmul(attention_weights, V)

        context_vector = tf.reduce_mean(context, axis=1)
        activation_weights = self.activation_weight_dense(context_vector)

        activation_weights_expanded = tf.expand_dims(
            tf.expand_dims(activation_weights, axis=1), axis=2
        )

        acts = tf.stack([f(x) for f in self.activations], axis=-1)
        weighted_activations = acts * activation_weights_expanded
        combined = tf.reduce_sum(weighted_activations, axis=-1)

        if not self.return_sequences:
            combined = tf.reduce_mean(combined, axis=1)

        if input_is_2d and self.return_sequences:
            combined = tf.squeeze(combined, axis=1)

        activation_sum = tf.reduce_sum(activation_weights, axis=0)
        self.activation_counts.assign_add(activation_sum)

        return combined

    def compute_output_shape(self, input_shape):
        if self.return_sequences:
            return (input_shape[0], input_shape[1], self.units)
        return (input_shape[0], self.units)

    def get_config(self):
        config = super().get_config()
        config.update({
            "units": self.units,
            "activations": self.activation_names,
            "return_sequences": self.return_sequences,
        })
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


# ---------------------------------------------------------------------------
# MAXLayerWithSelfAttentionV2 - GSM with channel-wise attention mode
# ---------------------------------------------------------------------------
@tf.keras.utils.register_keras_serializable()
class MAXLayerWithSelfAttentionV2(Layer):
    """Mixture-of-Activation layer with adaptive self-attention.

    Args:
        units: Internal projection dimension and Q/K/V dimension.
        activations: List of activation names to mix.
        return_sequences: Only relevant when attention_mode=="temporal".
            True keeps temporal dimension, False averages over it.
        attention_mode: "temporal" or "channel".
            "temporal" - attention over timesteps (original behavior).
            "channel"  - attention over features; meaningful for non-sequential data.
    """

    def __init__(self, units=32, activations=None, return_sequences=True,
                 attention_mode="channel", **kwargs):
        super().__init__(**kwargs)

        if attention_mode not in {"temporal", "channel"}:
            raise ValueError("attention_mode must be 'temporal' or 'channel'")

        self.units = units
        self.return_sequences = return_sequences
        self.attention_mode = attention_mode

        activations = activations or ["relu", "swish", "gelu"]
        self.activations, self.activation_names = _resolve_activations(activations)

        self.q_dense = Dense(units, use_bias=False)
        self.k_dense = Dense(units, use_bias=False)
        self.v_dense = Dense(units, use_bias=False)
        self.activation_weight_dense = Dense(
            len(self.activations), activation="softmax"
        )

    def build(self, input_shape):
        self.input_dim = input_shape[-1]
        n_act = len(self.activations)

        self.w = self.add_weight(
            name="kernel",
            shape=(self.input_dim, self.units),
            initializer="glorot_uniform",
            trainable=True,
        )
        self.b = self.add_weight(
            name="bias",
            shape=(self.units,),
            initializer="zeros",
            trainable=True,
        )
        self.activation_counts = self.add_weight(
            name="activation_counts",
            shape=(n_act,),
            initializer="zeros",
            trainable=False,
        )

        if self.attention_mode == "channel":
            dense_input_shape = (None, 1)
        else:
            dense_input_shape = (None, self.units)

        self.q_dense.build(dense_input_shape)
        self.k_dense.build(dense_input_shape)
        self.v_dense.build(dense_input_shape)
        self.activation_weight_dense.build((None, self.units))

        super().build(input_shape)

    def call(self, inputs):
        input_is_2d = (len(inputs.shape) == 2)
        if input_is_2d:
            inputs = tf.expand_dims(inputs, 1)

        x = tf.matmul(inputs, self.w) + self.b  # (batch, T, units)

        # Choose attention axis
        if self.attention_mode == "temporal":
            x_attn = x
        elif self.attention_mode == "channel":
            x_collapsed = tf.reduce_mean(x, axis=1)       # (batch, units)
            x_attn = tf.expand_dims(x_collapsed, 1)       # (batch, 1, units)
            x_attn = tf.transpose(x_attn, [0, 2, 1])      # (batch, units, 1)

        # Self-attention
        Q = self.q_dense(x_attn)
        K = self.k_dense(x_attn)
        V = self.v_dense(x_attn)

        d_k = tf.cast(self.units, tf.float32)
        scores = tf.matmul(Q, K, transpose_b=True) / tf.sqrt(d_k)
        attn = tf.nn.softmax(scores, axis=-1)
        context = tf.matmul(attn, V)

        context_vector = tf.reduce_mean(context, axis=1)  # (batch, units)
        pi = self.activation_weight_dense(context_vector)  # (batch, k)

        # Mix activations
        if self.attention_mode == "temporal":
            acts = tf.stack([f(x) for f in self.activations], axis=-1)
            pi_exp = tf.expand_dims(tf.expand_dims(pi, 1), 2)
            combined = tf.reduce_sum(acts * pi_exp, axis=-1)

            if not self.return_sequences:
                combined = tf.reduce_mean(combined, axis=1)

            if input_is_2d and self.return_sequences:
                combined = tf.squeeze(combined, axis=1)
        else:  # channel
            acts = tf.stack([f(context_vector) for f in self.activations], axis=-1)
            pi_exp = tf.expand_dims(pi, 1)
            combined = tf.reduce_sum(acts * pi_exp, axis=-1)

        self.activation_counts.assign_add(tf.reduce_sum(pi, axis=0))

        return combined

    def compute_output_shape(self, input_shape):
        if self.attention_mode == "temporal" and self.return_sequences:
            return (input_shape[0], input_shape[1], self.units)
        return (input_shape[0], self.units)

    def get_config(self):
        config = super().get_config()
        config.update({
            "units": self.units,
            "activations": self.activation_names,
            "return_sequences": self.return_sequences,
            "attention_mode": self.attention_mode,
        })
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)
