import torch
import blaze as bl


def test_get_parameter():
    def forward(x):
        scale = bl.get_parameter("scale", (1,), init_fn=torch.ones)
        bias = bl.get_parameter("bias", (1,), init_fn=torch.zeros)
        return x * scale + bias

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

    params = list(model.parameters())
    assert len(params) == 2

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


def test_get_parameter_reused():
    """get_parameter returns the same tensor on subsequent forward calls."""
    captured = []

    def forward(x):
        scale = bl.get_parameter("scale", (x.shape[-1],), init_fn=torch.ones)
        captured.append(scale)
        return x * scale

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

    captured.clear()
    model(torch.randn(2, 4))
    model(torch.randn(2, 4))

    assert len(captured) == 2
    assert captured[0].data_ptr() == captured[1].data_ptr()


def test_get_state_basic():
    def forward(x):
        buf = bl.get_state("running_mean", (x.shape[-1],), init_fn=torch.zeros)
        return x + buf

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

    assert "running_mean" in model._registry
    assert len(list(model.parameters())) == 0

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


def test_get_state_not_trainable():
    def forward(x):
        bl.get_state("buf", (4,))
        return x

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

    params = list(model.parameters())
    buffers = list(model.buffers())
    assert len(params) == 0
    assert len(buffers) == 1


def test_get_state_init_fn():
    def forward(x):
        offset = bl.get_state("offset", (3,), init_fn=torch.ones)
        return x + offset

    model = bl.transform(forward)
    model.init(torch.zeros(1, 3))

    x = torch.zeros(1, 3)
    out = model(x)
    torch.testing.assert_close(out, torch.ones(1, 3))


def test_get_state_scoped():
    class Counter(bl.Module):
        def __call__(self, x):
            bl.get_state("count", (1,), init_fn=torch.zeros)
            return x

    def forward(x):
        return Counter()(x)

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

    assert "counter.count" in model._registry


def test_set_state():
    def forward(x):
        count = bl.get_state("count", (1,), init_fn=torch.zeros)
        bl.set_state("count", count + 1)
        return x + count

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

    x = torch.zeros(1, 3)
    out1 = model(x)
    out2 = model(x)
    # Each call increments count by 1, so out2 - out1 == 1
    torch.testing.assert_close(out2 - out1, torch.ones(1, 3))


def test_set_state_without_get_state():
    def forward(x):
        bl.set_state("total", x.sum().unsqueeze(0))
        return x

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

    model(torch.ones(1, 3))
    assert model._registry["total"].buf.item() == 3.0

    model(torch.full((1, 3), 2.0))
    assert model._registry["total"].buf.item() == 6.0
