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


def test_demo_example():
    def forward(x, in_size=10, out_size=20):
        linear = bl.Linear(in_size, out_size)
        return linear(x)

    model = bl.transform(forward, in_size=10, out_size=20)
    model.init(torch.randn(5, 10))
    output = model(torch.randn(5, 10))
    assert output.shape == torch.Size([5, 20])


def test_multi_layer_naming():
    """Two Linear calls produce 'linear' and 'linear_1'."""

    def forward(x):
        x = bl.Linear(10, 32)(x)
        x = bl.ReLU()(x)
        x = bl.Linear(32, 5)(x)
        return x

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

    keys = sorted(model._registry.keys())
    assert keys == ["linear", "linear_1", "relu"]

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


def test_nested_module():
    class Block(bl.Module):
        def __call__(self, x, dim=32):
            x = bl.Linear(x.shape[-1], dim)(x)
            x = bl.ReLU()(x)
            return x

    def forward(x):
        b1 = Block()
        b2 = Block()
        x = b1(x, dim=64)
        x = b2(x, dim=10)
        return x

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

    keys = sorted(model._registry.keys())
    assert "block.linear" in keys
    assert "block.relu" in keys
    assert "block_1.linear" in keys
    assert "block_1.relu" in keys

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


def test_deeply_nested():
    class Inner(bl.Module):
        def __call__(self, x):
            return bl.Linear(x.shape[-1], x.shape[-1])(x)

    class Outer(bl.Module):
        def __call__(self, x):
            x = Inner()(x)
            x = Inner()(x)
            return x

    def forward(x):
        return Outer()(x)

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

    keys = sorted(model._registry.keys())
    assert "outer.inner.linear" in keys
    assert "outer.inner_1.linear" in keys

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


def test_loop():
    def forward(x):
        for _ in range(3):
            x = bl.Linear(10, 10)(x)
            x = bl.ReLU()(x)
        return x

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

    keys = sorted(model._registry.keys())
    assert "linear" in keys
    assert "linear_1" in keys
    assert "linear_2" in keys
    assert len([k for k in keys if k.startswith("linear")]) == 3

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


def test_weights_reused():
    """Verify that compile creates weights and forward reuses them (not new ones)."""

    def forward(x):
        return bl.Linear(10, 5)(x)

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

    weight_before = model._registry["linear"].weight.data.clone()

    x = torch.randn(1, 10)
    out = model(x)
    expected = torch.nn.functional.linear(x, weight_before, model._registry["linear"].bias)
    torch.testing.assert_close(out, expected)


def test_custom_name():
    def forward(x):
        x = bl.Linear(10, 10, name="encoder")(x)
        x = bl.Linear(10, 5, name="decoder")(x)
        return x

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

    keys = sorted(model._registry.keys())
    assert keys == ["decoder", "encoder"]


def test_error_not_compiled():
    def forward(x):
        return bl.Linear(10, 5)(x)

    model = bl.transform(forward)
    with pytest.raises(RuntimeError, match="not compiled"):
        model(torch.randn(1, 10))


def test_error_outside_transform():
    with pytest.raises(RuntimeError, match="No active blaze frame"):
        bl.Linear(10, 5)


def test_multiple_inputs_outputs():
    """Test that the model can handle multiple inputs and outputs."""

    def forward(x, y):
        x = bl.Linear(10, 20)(x)
        y = bl.Linear(10, 20)(y)
        z = x + y
        z = bl.ReLU()(z)
        w = bl.Linear(20, 5)(z)
        v = bl.Linear(20, 5)(z)
        return w, v

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

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


def test_multithreading():
    """Test that multiple threads can use the model without interfering."""

    def forward(x):
        x = bl.Linear(10, 20)(x)
        x = bl.ReLU()(x)
        x = bl.Linear(20, 5)(x)
        return x

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

    def thread_fn():
        for _ in range(10):
            out = model(torch.randn(2, 10))
            assert out.shape == (2, 5)

    import threading
    threads = [threading.Thread(target=thread_fn) for _ in range(4)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
