import pytest
import torch

from mini_transformer.modules.encoder import EncoderLayer, TransformerEncoder


def _dev():
    return (
        torch.device(f"cuda:{torch.cuda.current_device()}")
        if torch.cuda.is_available()
        else torch.device("cpu")
    )


# -------- ctor validations --------


def test_encoder_layer_param_checks():
    with pytest.raises(TypeError):
        EncoderLayer("32", 4, 64, 0.1)
    with pytest.raises(TypeError):
        EncoderLayer(32, "4", 64, 0.1)
    with pytest.raises(TypeError):
        EncoderLayer(32, 4, "64", 0.1)
    with pytest.raises(TypeError):
        EncoderLayer(32, 4, 64, "0.1")
    with pytest.raises(TypeError):
        EncoderLayer(32, 4, 64, 0.1, layer_norm_style=123)

    with pytest.raises(ValueError):
        EncoderLayer(0, 4, 64, 0.1)
    with pytest.raises(ValueError):
        EncoderLayer(32, 0, 64, 0.1)
    with pytest.raises(ValueError):
        EncoderLayer(32, 4, 0, 0.1)
    with pytest.raises(ValueError):
        EncoderLayer(32, 4, 64, 1.0)  # upper bound excluded
    with pytest.raises(ValueError):
        EncoderLayer(32, 4, 64, 0.1, layer_norm_style="middle")


def test_transformer_encoder_layers_count_checks():
    with pytest.raises(TypeError):
        TransformerEncoder(32, 4, 64, "2", 0.1)
    with pytest.raises(ValueError):
        TransformerEncoder(32, 4, 64, 0, 0.1)


# -------- forward path --------


@pytest.mark.parametrize("B,S,D,H,FF,L", [(2, 5, 24, 3, 48, 2)])
def test_encoder_forward_happy_path(B, S, D, H, FF, L):
    device = _dev()
    enc = TransformerEncoder(D, H, FF, L, 0.1).to(device)

    x = torch.randn(B, S, D, device=device)
    heads = enc.layers[0].attention_layer.num_heads
    src_pad = torch.zeros(B, heads, 1, S, dtype=torch.bool, device=device)

    out = enc(x, src_pad)
    assert out.shape == (B, S, D)
    assert out.device == device


def test_transformer_encoder_pre_norm_layers_flag():
    enc = TransformerEncoder(24, 3, 48, 2, 0.1, layer_norm_style="pre")
    assert all(layer.pre_norm for layer in enc.layers)


def test_encoder_pre_norm_forward_matches_shapes():
    device = _dev()
    layer = EncoderLayer(24, 3, 48, 0.1, layer_norm_style="pre").to(device)

    x = torch.randn(2, 5, 24, device=device)
    mask = torch.zeros(2, 3, 1, 5, dtype=torch.bool, device=device)
    out = layer(x, mask)
    assert out.shape == x.shape
    assert layer.pre_norm is True


def test_encoder_forward_input_checks_and_message_format():
    layer = EncoderLayer(24, 3, 48, 0.1)

    with pytest.raises(TypeError):
        layer("not a tensor", None)

    with pytest.raises(ValueError) as ei:
        layer(torch.randn(2, 3, 4, 5), None)  # rank 4
    assert "x must be a 3D torch.Tensor of shape (B, S, D)" in str(ei.value)