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


def test_init_fn_layer():
    def forward(x):
        x = bl.Linear(10, 20, init_fn=lambda m: nn.init.zeros_(m.weight))(x)
        return x

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

    weight = model._registry["linear"].weight
    torch.testing.assert_close(weight, torch.zeros_like(weight))

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


def test_init_fn():
    class Block(bl.Module):
        def __init__(self, dim, init_fn=None):
            super().__init__(init_fn=init_fn)
            self.dim = dim

        def __call__(self, x):
            return bl.Linear(x.shape[-1], self.dim)(x)

    def forward(x):
        return Block(dim=5, init_fn=lambda m: setattr(m, 'tag', 'initialized'))(x)

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

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


def test_init_fn_multiple_modules():
    inited = []

    def track(name):
        def fn(m):
            inited.append(name)
        return fn

    def forward(x):
        x = bl.Linear(10, 20, init_fn=track("linear"))(x)
        x = bl.Conv1d(20, 8, kernel_size=3, padding=1, init_fn=track("conv"))(x.unsqueeze(-1)).squeeze(-1)
        x = bl.BatchNorm1d(8, init_fn=track("bn"))(x)
        x = bl.LayerNorm(8, init_fn=track("ln"))(x)
        x = bl.Linear(8, 5, init_fn=track("linear2"))(x)
        return x

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

    assert inited == ["linear", "conv", "bn", "ln", "linear2"]

    # init_fn not called again during apply
    inited.clear()
    out = model(torch.randn(2, 10))
    assert inited == []
    assert out.shape == (2, 5)


def test_init_fn_not_called_during_apply():
    call_count = [0]

    def my_init(m):
        call_count[0] += 1

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

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))
    assert call_count[0] == 1

    model(torch.randn(2, 10))
    model(torch.randn(2, 10))
    assert call_count[0] == 1  # not called again during APPLY
