import pytest
import torch

from mini_transformer.modules.feedforward import FeedForwardLayer

# =======================
# Constructor checks
# =======================


def test_ctor_type_checks():
    with pytest.raises(TypeError):
        FeedForwardLayer("64", 256, 0.1)  # d_model
    with pytest.raises(TypeError):
        FeedForwardLayer(64, "256", 0.1)  # d_ff
    with pytest.raises(TypeError):
        FeedForwardLayer(64, 256, "0.1")  # dropout_rate


def test_ctor_value_checks():
    with pytest.raises(ValueError):
        FeedForwardLayer(0, 256, 0.1)
    with pytest.raises(ValueError):
        FeedForwardLayer(64, 0, 0.1)
    with pytest.raises(ValueError):
        FeedForwardLayer(64, 256, -1.0)
    with pytest.raises(ValueError):
        FeedForwardLayer(64, 256, 1.0)  # upper bound excluded


def test_ctor_happy_path_defaults_and_attrs():
    ffn = FeedForwardLayer(64, 256)  # default dropout=0.1
    assert ffn.d_model == 64 and ffn.d_ff == 256
    assert isinstance(ffn.dropout, torch.nn.Dropout)
    assert abs(ffn.dropout.p - 0.1) < 1e-9


# =======================
# Forward: input validation
# =======================


def test_forward_type_and_rank_checks():
    ffn = FeedForwardLayer(32, 64, 0.1)
    with pytest.raises(TypeError):
        ffn("not a tensor")
    with pytest.raises(ValueError):
        ffn(torch.randn(5, 32))  # rank 2
    with pytest.raises(ValueError):
        ffn(torch.randn(2, 3, 16))  # D mismatch (expects 32)


# =======================
# Forward: shapes, device, dtype, zero-length
# =======================


@pytest.mark.parametrize("B,S,D", [(1, 1, 16), (2, 7, 32), (3, 0, 24)])
def test_forward_shapes_device_dtype_and_zero_len(B, S, D):
    device = (
        torch.device(f"cuda:{torch.cuda.current_device()}")
        if torch.cuda.is_available()
        else torch.device("cpu")
    )
    ffn = FeedForwardLayer(D, D * 2, 0.2).to(device)

    x = torch.randn(B, S, D, device=device, dtype=torch.float32)
    y = ffn(x)
    assert y.shape == (B, S, D)
    assert y.device == device
    assert y.dtype == x.dtype


# =======================
# Forward: gradient flow
# =======================


def test_gradients_flow():
    D, H = 32, 64
    ffn = FeedForwardLayer(D, H, 0.1)
    x = torch.randn(2, 5, D, requires_grad=True)
    y = ffn(x)
    loss = y.pow(2).mean()
    loss.backward()
    assert x.grad is not None and torch.isfinite(x.grad).all()


# =======================
# Dropout behavior
# =======================


def test_dropout_zero_equals_no_dropout_path():
    D, H = 16, 32
    # Build two modules: one with p=0.0, one with p>0 but switched to eval()
    ffn0 = FeedForwardLayer(D, H, 0.0)
    ffnx = FeedForwardLayer(D, H, 0.5).eval()  # eval disables dropout

    ffnx.load_state_dict(ffn0.state_dict(), strict=False)

    x = torch.randn(2, 4, D)
    y0 = ffn0(x)
    yx = ffnx(x)
    # With dropout disabled both should be equal (same weights distrib not identical, but same computation tree)
    assert torch.allclose(y0, yx, atol=1e-6, rtol=1e-6)


def test_dropout_training_changes_output_vs_eval():
    D, H = 32, 64
    ffn = FeedForwardLayer(D, H, 0.5)
    x = torch.randn(3, 6, D)

    torch.manual_seed(123)
    ffn.train()
    y_train = ffn(x)

    torch.manual_seed(123)
    ffn.eval()
    y_eval = ffn(x)

    # Same seed but eval disables dropout -> outputs should differ
    assert not torch.allclose(y_train, y_eval)


def test_dropout_is_noop_on_zero_length():
    D, H = 32, 64
    ffn = FeedForwardLayer(D, H, 0.7).train()
    x = torch.randn(2, 0, D)  # zero-length sequence
    y = ffn(x)
    assert y.shape == (2, 0, D)  # no crash; shape preserved