import pytest
import numpy as np
import tensorflow as tf
from mombai.layers.multi_activation import (
    ACTIVATION_MAP,
    _resolve_activations,
    WAFLayer,
    WAFLayerNorm,
    WAFLayerNormV2,
    MAXLayerWithAttention,
    MAXLayerWithSelfAttention,
    MAXLayerWithSelfAttentionV2,
)


@pytest.fixture
def input_2d():
    return tf.constant(np.random.randn(4, 16).astype(np.float32))


@pytest.fixture
def input_3d():
    return tf.constant(np.random.randn(4, 5, 16).astype(np.float32))


class TestResolveActivations:
    def test_string_activations(self):
        funcs, names = _resolve_activations(["relu", "swish"])
        assert len(funcs) == 2
        assert names == ["relu", "swish"]

    def test_function_activations(self):
        funcs, names = _resolve_activations([tf.nn.relu])
        assert len(funcs) == 1
        assert names == ["relu"]

    def test_unknown_activation_raises(self):
        with pytest.raises(ValueError, match="Unknown activation"):
            _resolve_activations(["nonexistent_activation"])

    def test_mixed_activations(self):
        funcs, names = _resolve_activations(["relu", tf.nn.tanh])
        assert len(funcs) == 2
        assert names[0] == "relu"
        assert names[1] == "tanh"


class TestActivationMap:
    def test_contains_standard_activations(self):
        expected = ["relu", "sigmoid", "tanh", "softplus", "elu", "selu",
                    "swish", "gelu", "leaky_relu"]
        for name in expected:
            assert name in ACTIVATION_MAP

    def test_crelu_not_present(self):
        assert "crelu" not in ACTIVATION_MAP


class TestWAFLayer:
    def test_default_activations(self):
        layer = WAFLayer(units=8)
        assert layer.activation_names == ["relu", "swish"]

    def test_custom_activations(self):
        layer = WAFLayer(units=8, activations=["tanh", "sigmoid", "elu"])
        assert layer.activation_names == ["tanh", "sigmoid", "elu"]

    def test_output_shape_2d(self, input_2d):
        layer = WAFLayer(units=8, activations=["relu", "tanh"])
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_compressor_sum(self, input_2d):
        layer = WAFLayer(units=8, activations=["relu", "tanh"], compressor="sum")
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_compressor_avg(self, input_2d):
        layer = WAFLayer(units=8, activations=["relu", "tanh"], compressor="avg")
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_invalid_compressor(self, input_2d):
        layer = WAFLayer(units=8, activations=["relu", "tanh"], compressor="max")
        with pytest.raises(ValueError, match="Invalid compressor"):
            layer(input_2d)

    def test_father_activation_enabled(self, input_2d):
        layer = WAFLayer(units=8, activations=["relu", "tanh"],
                         apply_father_activation=True)
        output = layer(input_2d)
        assert output.shape == (4, 8)
        assert hasattr(layer, "softingparam")

    def test_father_activation_disabled(self, input_2d):
        layer = WAFLayer(units=8, activations=["relu", "tanh"],
                         apply_father_activation=False)
        output = layer(input_2d)
        assert output.shape == (4, 8)
        assert not hasattr(layer, "softingparam")

    def test_compute_output_shape(self):
        layer = WAFLayer(units=8)
        assert layer.compute_output_shape((None, 16)) == (None, 8)

    def test_get_config(self):
        layer = WAFLayer(units=8, activations=["relu", "gelu"],
                         compressor="avg", apply_father_activation=True)
        config = layer.get_config()
        assert config["units"] == 8
        assert config["activations"] == ["relu", "gelu"]
        assert config["compressor"] == "avg"
        assert config["apply_father_activation"] is True

    def test_from_config(self, input_2d):
        layer = WAFLayer(units=8, activations=["relu", "gelu"], compressor="avg")
        layer(input_2d)
        config = layer.get_config()
        restored = WAFLayer.from_config(config)
        output = restored(input_2d)
        assert output.shape == (4, 8)

    def test_serialization_roundtrip(self, input_2d):
        layer = WAFLayer(units=8, activations=["sigmoid", "tanh"],
                         apply_father_activation=True)
        layer(input_2d)
        config = layer.get_config()
        restored = WAFLayer.from_config(config)
        assert restored.activation_names == ["sigmoid", "tanh"]
        assert restored.apply_father_activation is True


class TestWAFLayerNorm:
    def test_output_shape(self, input_2d):
        layer = WAFLayerNorm(units=8, activations=["relu", "tanh"])
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_compressor_sum(self, input_2d):
        layer = WAFLayerNorm(units=8, activations=["relu", "tanh"], compressor="sum")
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_compressor_avg(self, input_2d):
        layer = WAFLayerNorm(units=8, activations=["relu", "tanh"], compressor="avg")
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_father_activation(self, input_2d):
        layer = WAFLayerNorm(units=8, activations=["relu", "tanh"],
                             apply_father_activation=True)
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_from_config(self, input_2d):
        layer = WAFLayerNorm(units=8, activations=["relu", "gelu"])
        layer(input_2d)
        config = layer.get_config()
        restored = WAFLayerNorm.from_config(config)
        output = restored(input_2d)
        assert output.shape == (4, 8)


class TestWAFLayerNormV2:
    def test_output_shape(self, input_2d):
        layer = WAFLayerNormV2(units=8, activations=["relu", "tanh"])
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_compressor_sum(self, input_2d):
        layer = WAFLayerNormV2(units=8, activations=["relu", "exponential"],
                               compressor="sum")
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_compressor_avg(self, input_2d):
        layer = WAFLayerNormV2(units=8, activations=["relu", "exponential"],
                               compressor="avg")
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_father_activation(self, input_2d):
        layer = WAFLayerNormV2(units=8, activations=["relu", "tanh"],
                               apply_father_activation=True)
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_normalizes_per_branch(self, input_2d):
        layer = WAFLayerNormV2(units=8, activations=["relu", "exponential"])
        output = layer(input_2d)
        assert not tf.reduce_any(tf.math.is_nan(output))

    def test_from_config(self, input_2d):
        layer = WAFLayerNormV2(units=8, activations=["relu", "gelu"])
        layer(input_2d)
        config = layer.get_config()
        restored = WAFLayerNormV2.from_config(config)
        output = restored(input_2d)
        assert output.shape == (4, 8)


class TestMAXLayerWithAttention:
    def test_output_shape_2d(self, input_2d):
        layer = MAXLayerWithAttention(units=8, activations=["relu", "tanh"])
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_output_shape_3d_return_sequences(self, input_3d):
        layer = MAXLayerWithAttention(units=8, activations=["relu", "tanh"],
                                      return_sequences=True)
        output = layer(input_3d)
        assert output.shape == (4, 5, 8)

    def test_output_shape_3d_no_return_sequences(self, input_3d):
        layer = MAXLayerWithAttention(units=8, activations=["relu", "tanh"],
                                      return_sequences=False)
        output = layer(input_3d)
        assert output.shape == (4, 8)

    def test_activation_counts_tracked(self, input_2d):
        layer = MAXLayerWithAttention(units=8, activations=["relu", "tanh"])
        layer(input_2d)
        counts = layer.activation_counts.numpy()
        assert counts.shape == (2,)
        assert np.sum(counts) > 0

    def test_default_activations(self):
        layer = MAXLayerWithAttention(units=8)
        assert layer.activation_names == ["relu", "swish"]

    def test_compute_output_shape_with_sequences(self):
        layer = MAXLayerWithAttention(units=8, return_sequences=True)
        assert layer.compute_output_shape((None, 5, 16)) == (None, 5, 8)

    def test_compute_output_shape_without_sequences(self):
        layer = MAXLayerWithAttention(units=8, return_sequences=False)
        assert layer.compute_output_shape((None, 5, 16)) == (None, 8)

    def test_get_config(self):
        layer = MAXLayerWithAttention(units=8, activations=["relu", "gelu"],
                                      return_sequences=False)
        config = layer.get_config()
        assert config["units"] == 8
        assert config["activations"] == ["relu", "gelu"]
        assert config["return_sequences"] is False

    def test_from_config(self, input_2d):
        layer = MAXLayerWithAttention(units=8, activations=["relu", "gelu"])
        layer(input_2d)
        config = layer.get_config()
        restored = MAXLayerWithAttention.from_config(config)
        output = restored(input_2d)
        assert output.shape == (4, 8)


class TestMAXLayerWithSelfAttention:
    def test_output_shape_2d(self, input_2d):
        layer = MAXLayerWithSelfAttention(units=8, activations=["relu", "tanh"])
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_output_shape_3d_return_sequences(self, input_3d):
        layer = MAXLayerWithSelfAttention(units=8, activations=["relu", "tanh"],
                                          return_sequences=True)
        output = layer(input_3d)
        assert output.shape == (4, 5, 8)

    def test_output_shape_3d_no_return_sequences(self, input_3d):
        layer = MAXLayerWithSelfAttention(units=8, activations=["relu", "tanh"],
                                          return_sequences=False)
        output = layer(input_3d)
        assert output.shape == (4, 8)

    def test_activation_counts_tracked(self, input_2d):
        layer = MAXLayerWithSelfAttention(units=8, activations=["relu", "tanh"])
        layer(input_2d)
        counts = layer.activation_counts.numpy()
        assert counts.shape == (2,)
        assert np.sum(counts) > 0

    def test_compute_output_shape_with_sequences(self):
        layer = MAXLayerWithSelfAttention(units=8, return_sequences=True)
        assert layer.compute_output_shape((None, 5, 16)) == (None, 5, 8)

    def test_compute_output_shape_without_sequences(self):
        layer = MAXLayerWithSelfAttention(units=8, return_sequences=False)
        assert layer.compute_output_shape((None, 5, 16)) == (None, 8)

    def test_get_config(self):
        layer = MAXLayerWithSelfAttention(units=8, activations=["relu", "gelu"],
                                          return_sequences=False)
        config = layer.get_config()
        assert config["units"] == 8
        assert config["return_sequences"] is False

    def test_from_config(self, input_2d):
        layer = MAXLayerWithSelfAttention(units=8, activations=["relu", "gelu"])
        layer(input_2d)
        config = layer.get_config()
        restored = MAXLayerWithSelfAttention.from_config(config)
        output = restored(input_2d)
        assert output.shape == (4, 8)


class TestMAXLayerWithSelfAttentionV2:
    def test_channel_mode_2d(self, input_2d):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, activations=["relu", "tanh"], attention_mode="channel")
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_temporal_mode_2d(self, input_2d):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, activations=["relu", "tanh"], attention_mode="temporal")
        output = layer(input_2d)
        assert output.shape == (4, 8)

    def test_temporal_mode_3d_return_sequences(self, input_3d):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, activations=["relu", "tanh"],
            attention_mode="temporal", return_sequences=True)
        output = layer(input_3d)
        assert output.shape == (4, 5, 8)

    def test_temporal_mode_3d_no_return_sequences(self, input_3d):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, activations=["relu", "tanh"],
            attention_mode="temporal", return_sequences=False)
        output = layer(input_3d)
        assert output.shape == (4, 8)

    def test_channel_mode_3d(self, input_3d):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, activations=["relu", "tanh"], attention_mode="channel")
        output = layer(input_3d)
        assert output.shape == (4, 8)

    def test_invalid_attention_mode(self):
        with pytest.raises(ValueError, match="attention_mode"):
            MAXLayerWithSelfAttentionV2(units=8, attention_mode="invalid")

    def test_default_activations(self):
        layer = MAXLayerWithSelfAttentionV2(units=8)
        assert layer.activation_names == ["relu", "swish", "gelu"]

    def test_activation_counts_tracked(self, input_2d):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, activations=["relu", "tanh"], attention_mode="channel")
        layer(input_2d)
        counts = layer.activation_counts.numpy()
        assert counts.shape == (2,)
        assert np.sum(counts) > 0

    def test_compute_output_shape_channel(self):
        layer = MAXLayerWithSelfAttentionV2(units=8, attention_mode="channel")
        assert layer.compute_output_shape((None, 16)) == (None, 8)

    def test_compute_output_shape_temporal_with_seq(self):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, attention_mode="temporal", return_sequences=True)
        assert layer.compute_output_shape((None, 5, 16)) == (None, 5, 8)

    def test_compute_output_shape_temporal_no_seq(self):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, attention_mode="temporal", return_sequences=False)
        assert layer.compute_output_shape((None, 5, 16)) == (None, 8)

    def test_get_config(self):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, activations=["relu", "gelu"],
            attention_mode="channel", return_sequences=False)
        config = layer.get_config()
        assert config["units"] == 8
        assert config["attention_mode"] == "channel"
        assert config["return_sequences"] is False

    def test_from_config_channel(self, input_2d):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, activations=["relu", "gelu"], attention_mode="channel")
        layer(input_2d)
        config = layer.get_config()
        restored = MAXLayerWithSelfAttentionV2.from_config(config)
        output = restored(input_2d)
        assert output.shape == (4, 8)

    def test_from_config_temporal(self, input_3d):
        layer = MAXLayerWithSelfAttentionV2(
            units=8, activations=["relu", "gelu"], attention_mode="temporal")
        layer(input_3d)
        config = layer.get_config()
        restored = MAXLayerWithSelfAttentionV2.from_config(config)
        output = restored(input_3d)
        assert output.shape == (4, 5, 8)


class TestKerasIntegration:
    def test_waf_in_model(self, input_2d):
        model = tf.keras.Sequential([
            tf.keras.layers.InputLayer(input_shape=(16,)),
            WAFLayer(units=8, activations=["relu", "tanh"]),
            tf.keras.layers.Dense(1, activation="sigmoid"),
        ])
        output = model(input_2d)
        assert output.shape == (4, 1)

    def test_waf_norm_v2_in_model(self, input_2d):
        model = tf.keras.Sequential([
            tf.keras.layers.InputLayer(input_shape=(16,)),
            WAFLayerNormV2(units=8, activations=["relu", "tanh"]),
            tf.keras.layers.Dense(1, activation="sigmoid"),
        ])
        output = model(input_2d)
        assert output.shape == (4, 1)

    def test_max_attention_in_model(self, input_2d):
        model = tf.keras.Sequential([
            tf.keras.layers.InputLayer(input_shape=(16,)),
            MAXLayerWithAttention(units=8, activations=["relu", "tanh"]),
            tf.keras.layers.Dense(1, activation="sigmoid"),
        ])
        output = model(input_2d)
        assert output.shape == (4, 1)

    def test_max_self_attention_v2_in_model(self, input_2d):
        model = tf.keras.Sequential([
            tf.keras.layers.InputLayer(input_shape=(16,)),
            MAXLayerWithSelfAttentionV2(units=8, activations=["relu", "tanh"],
                                         attention_mode="channel"),
            tf.keras.layers.Dense(1, activation="sigmoid"),
        ])
        output = model(input_2d)
        assert output.shape == (4, 1)

    def test_model_trainable(self, input_2d):
        model = tf.keras.Sequential([
            tf.keras.layers.InputLayer(input_shape=(16,)),
            WAFLayerNormV2(units=8, activations=["relu", "tanh"]),
            tf.keras.layers.Dense(1, activation="sigmoid"),
        ])
        y = tf.constant(np.random.randint(0, 2, (4, 1)).astype(np.float32))
        model.compile(optimizer="adam", loss="binary_crossentropy")
        loss_before = model.evaluate(input_2d, y, verbose=0)
        model.fit(input_2d, y, epochs=5, verbose=0)
        loss_after = model.evaluate(input_2d, y, verbose=0)
        assert loss_before != loss_after
