import torch
import torch.nn as nn
import blaze as bl


def test_TransformerEncoderLayer():
    """Test that TransformerEncoderLayer can be used."""

    def forward(x):
        x = bl.TransformerEncoderLayer(d_model=10, nhead=2)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 5, 10))

    out = model(torch.randn(2, 5, 10))
    assert out.shape == (2, 5, 10)


def test_TransformerEncoder():
    """Test that TransformerEncoder can be used."""

    def forward(x):
        layer = bl.TransformerEncoderLayer(d_model=10, nhead=2)
        x = bl.TransformerEncoder(layer, num_layers=2)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 5, 10))

    out = model(torch.randn(2, 5, 10))
    assert out.shape == (2, 5, 10)


def test_TransformerDecoder():
    """Test that TransformerDecoder can be used."""

    def forward(x, memory):
        layer = bl.TransformerDecoderLayer(d_model=10, nhead=2)
        x = bl.TransformerDecoder(layer, num_layers=2)(x, memory)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 5, 10), torch.randn(2, 5, 10))

    out = model(torch.randn(2, 5, 10), torch.randn(2, 5, 10))
    assert out.shape == (2, 5, 10)


def test_subclass():
    class PatchedTransformerEncoderLayer(bl.TransformerEncoderLayer):
        def __init__(self, alpha, attn_dropout, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.attn_weights = None
            self.attn_dropout = attn_dropout
            self.alpha = alpha
            self.linear = bl.Linear(self.self_attn.embed_dim, self.self_attn.embed_dim)

        def _sa_block(
            self,
            x: torch.Tensor,
            attn_mask: torch.Tensor | None,
            key_padding_mask: torch.Tensor | None,
            is_causal: bool = False,
        ) -> torch.Tensor:
            if attn_mask is not None:
                attn_mask = torch.where(torch.rand_like(attn_mask, dtype=torch.float32) < self.attn_dropout * self.training, torch.zeros_like(attn_mask), attn_mask)
            x, attn_weights = self.self_attn(
                x,
                x,
                x,
                attn_mask=attn_mask,
                key_padding_mask=key_padding_mask,
                need_weights=True,
                is_causal=is_causal,
            )
            x = self.linear(x)
            self.attn_weights = bl.Sigmoid()(bl.Linear(attn_weights.shape[-1], attn_weights.shape[-1])(attn_weights) * self.alpha)
            return self.dropout1(x)

    def forward(x):
        layer = PatchedTransformerEncoderLayer(alpha=10, attn_dropout=0.5, d_model=10, nhead=2, batch_first=True)
        x = layer(x)
        attn_weights = layer.attn_weights
        return x, attn_weights, torch.tensor(layer.alpha)

    model = bl.transform(forward)
    model.init(torch.randn(2, 5, 10))
    out = model(torch.randn(2, 5, 10))
    assert out[0].shape == (2, 5, 10)
    assert out[1].shape == (2, 5, 5)
    assert out[2] == 10
    assert (out[1] >= 0).all() and (out[1] <= 1).all()


def test_wrap():
    def forward(x):
        x = bl.wrap(lambda: nn.Linear(10, 20))(x)
        x = bl.wrap(lambda: nn.Linear(20, 5))(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    assert sorted(model._registry.keys()) == ["wrapped", "wrapped_1"]
    out = model(torch.randn(2, 10))
    assert out.shape == (2, 5)


def test_wrap_nested():
    class Block(bl.Module):
        def __call__(self, x):
            return bl.wrap(lambda: nn.Sequential(nn.Linear(10, 10), nn.ReLU()))(x)

    def forward(x):
        return Block()(x)

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    assert "block.wrapped" in model._registry
    out = model(torch.randn(2, 10))
    assert out.shape == (2, 10)


def test_conv2d():
    def forward(x):
        x = bl.Conv2d(3, 16, kernel_size=3, padding=1)(x)
        x = bl.BatchNorm2d(16)(x)
        x = bl.ReLU()(x)
        x = bl.MaxPool2d(2)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 3, 32, 32))

    out = model(torch.randn(2, 3, 32, 32))
    assert out.shape == (2, 16, 16, 16)
