import torch
import blaze as bl


def test_if_else():
    """If/else with shape-dependent condition."""
    def forward(x):
        x = bl.Linear(10, 20)(x)
        if x.shape[0] > 1:
            x = bl.ReLU()(x)
        else:
            x = bl.GELU()(x)
        x = bl.Linear(20, 5)(x)
        return x

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

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


def test_loop_eager():
    """Fixed-count loop in eager mode."""
    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(2, 10))

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


def test_conditional_skip_connection():
    """Residual skip connection."""
    def forward(x):
        identity = x
        x = bl.Linear(10, 10)(x)
        x = bl.ReLU()(x)
        x = x + identity  # skip connection
        x = bl.Linear(10, 5)(x)
        return x

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

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


def test_loop_with_if():
    """Loop body with if/else based on loop index."""
    def forward(x):
        for i in range(4):
            x = bl.Linear(10, 10)(x)
            if i % 2 == 0:
                x = bl.ReLU()(x)
            else:
                x = bl.GELU()(x)
        return x

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

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


def test_nested_module_with_control_flow():
    """bl.Module with internal if/else."""
    class Block(bl.Module):
        def __call__(self, x, use_bn=False):
            x = bl.Linear(x.shape[-1], 32)(x)
            if use_bn:
                x = bl.BatchNorm1d(32)(x)
            x = bl.ReLU()(x)
            return x

    def forward(x):
        x = Block()(x, use_bn=True)
        x = Block()(x, use_bn=False)
        return bl.Linear(32, 5)(x)

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

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