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


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

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

    params = list(model.parameters())
    assert len(params) == 4  # 2 weights + 2 biases

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    out = model(torch.randn(2, 10))
    loss = out.sum()
    loss.backward()
    optimizer.step()


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

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

    sd = model.state_dict()
    assert any("weight" in k for k in sd)
    assert any("bias" in k for k in sd)

    # Load into a fresh model
    model2 = bl.transform(forward)
    model2.init(torch.randn(1, 10))
    model2.load_state_dict(sd)

    x = torch.randn(1, 10)
    torch.testing.assert_close(model(x), model2(x))


def test_load_state_dict_partial():
    """Load a subset of weights from a larger model into a smaller one."""
    class Encoder(bl.Module):
        def __call__(self, x):
            x = x.view(x.shape[0], -1)
            return bl.Linear(28, 10)(x)

    class Decoder(bl.Module):
        def __call__(self, z):
            return bl.Linear(10, 28)(z)

    def full_forward(x):
        z = Encoder()(x)
        return Decoder()(z)

    model = bl.transform(full_forward)
    model.init(torch.randn(2, 28))

    # Train one step so weights differ from init
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    loss = model(torch.randn(2, 28)).sum()
    loss.backward()
    optimizer.step()

    # Extract decoder weights (user strips prefix, common PyTorch pattern)
    decoder_sd = {k: v
                  for k, v in model.state_dict().items()
                  if k.startswith('decoder.')}

    # Build standalone decoder
    def decode_only(z):
        return Decoder()(z)

    decoder = bl.transform(decode_only)
    decoder.init(torch.randn(1, 10))

    # This should work — load_state_dict should handle the key mismatch
    decoder.load_state_dict(decoder_sd)

    # Verify weights match
    assert set(decoder.state_dict().keys()) == set(decoder_sd.keys())


def test_train_eval_propagation():
    def forward(x):
        x = bl.Linear(10, 10)(x)
        x = bl.Dropout(0.5)(x)
        return x

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

    model.eval()
    x = torch.randn(100, 10)
    out1 = model(x)
    out2 = model(x)
    torch.testing.assert_close(out1, out2)  # deterministic in eval

    model.train()
    # In train mode, dropout introduces randomness (usually)
    # Just check it runs without error
    _ = model(x)


def test_batchnorm():
    def forward(x):
        x = bl.Linear(10, 20)(x)
        x = bl.BatchNorm1d(20)(x)
        return x

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

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

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


def test_batchnorm_buffers_captured():
    def forward(x):
        return bl.BatchNorm1d(3)(x)

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

    # BatchNorm has running_mean and running_var buffers
    bn = model._registry["batch_norm1d"]
    assert hasattr(bn, "running_mean")
    assert hasattr(bn, "running_var")

    # Buffers should appear in model state_dict via the traced module
    sd = model.state_dict()
    buf_keys = [k for k in sd if "running" in k or "num_batches" in k]
    assert len(buf_keys) > 0, f"no buffer keys found in state_dict: {list(sd.keys())}"

    # Running stats should update during forward in train mode
    model.train()
    model(torch.ones(4, 3) * 5.0)
    assert not torch.equal(bn.running_mean, torch.zeros(3))


def test_parameter_names_clean():
    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))

    for name, _ in model.named_parameters():
        assert not name.startswith("_traced."), f"bad prefix: {name}"

    expected = {"linear.weight", "linear.bias", "linear_1.weight", "linear_1.bias"}
    assert set(model.state_dict().keys()) == expected


def test_to_device():
    def forward(x):
        x = bl.Linear(10, 20)(x)
        x = bl.BatchNorm1d(20)(x)
        return x

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

    model = model.to("meta")
    for _, p in model.named_parameters():
        assert p.device.type == "meta"
    for _, b in model.named_buffers():
        assert b.device.type == "meta"
    for mod in model._registry.values():
        for p in mod.parameters():
            assert p.device.type == "meta"


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

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))
    model = torch.jit.trace(model, torch.randn(2, 10))
    model = model.to("meta")

    for p in model.parameters():
        assert p.device.type == "meta"


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

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))
    model = torch.compile(model, backend="eager")

    model.to("meta")
    for _, p in model.named_parameters():
        assert p.device.type == "meta"
    for mod in model._registry.values():
        for p in mod.parameters():
            assert p.device.type == "meta"
