import torch, sys
sys.path.insert(0, 'src')
from gigatorch.ops import ConvBlockND, SAPETransformerBlockND, AttentionND

# Test ConvBlockND
for dim in [1, 2, 3]:
    C = 256
    spatial = [16] * dim
    x = torch.randn(2, C, *spatial)
    block = ConvBlockND(dim, C)
    y = block(x)
    assert y.shape == x.shape, f'ConvBlock dim={dim}: {y.shape} vs {x.shape}'
    print(f'ConvBlockND dim={dim}: {x.shape} -> {y.shape} OK')

# Test SAPETransformerBlockND
for dim in [1, 2, 3]:
    C = 256
    spatial = [8] * dim
    x = torch.randn(2, C, *spatial)
    block = SAPETransformerBlockND(dim, C, num_heads=8, head_dim=32)
    y = block(x)
    assert y.shape == x.shape, f'TransformerBlock dim={dim}: {y.shape} vs {x.shape}'
    inner = C + dim
    print(f'SAPETransformerBlockND dim={dim}: {x.shape} -> {y.shape} (inner={inner}) OK')

# Test interleaving: 256c 258t 256c (the user's example, dim=2)
print()
print('=== Interleaved Sequential (256c -> 258t -> 256c) ===')
model = torch.nn.Sequential(
    ConvBlockND(2, 256),
    ConvBlockND(2, 256),
    SAPETransformerBlockND(2, 256, num_heads=8),
    ConvBlockND(2, 256),
    SAPETransformerBlockND(2, 256, num_heads=8),
    ConvBlockND(2, 256),
)
x = torch.randn(2, 256, 16, 16)
y = model(x)
print(f'Input: {x.shape}')
print(f'Output: {y.shape}')
assert y.shape == x.shape

# Gradient flow
loss = y.sum()
loss.backward()
print('Gradient flow OK')

# Parameter count
total = sum(p.numel() for p in model.parameters())
print(f'Total params: {total:,}')
