import torch, sys
sys.path.insert(0, 'src')
from gigatorch.ops import SAPE, SAPEOverwrite, AttentionND, ChannelAffine, Residual

for dim in [1, 2, 3]:
    C, H = 32, 4
    spatial = [8] * dim
    x = torch.randn(2, C, *spatial)

    # SAPEOverwrite: last dim channels overwritten
    sape = SAPEOverwrite(dim)
    y = sape(x)
    assert y.shape == x.shape, f'SAPEOverwrite shape mismatch: {y.shape} vs {x.shape}'
    # Verify position channels are correct (center should be ~0 for even-sized dims)
    assert torch.allclose(y[:, -dim:], sape(torch.zeros_like(x))[:, -dim:]), 'Position channels should be independent of content'

    # AttentionND
    attn = AttentionND(dim, C, H)
    y = attn(x)
    assert y.shape == x.shape, f'AttentionND shape mismatch: {y.shape} vs {x.shape}'

    # Composition: SAPEOverwrite -> ChannelAffine -> AttentionND inside Residual
    block = torch.nn.Sequential(
        SAPEOverwrite(dim),
        Residual(torch.nn.Sequential(ChannelAffine(dim, C), AttentionND(dim, C, H)))
    )
    y = block(x)
    assert y.shape == x.shape, f'Block shape mismatch: {y.shape} vs {x.shape}'

    print(f'dim={dim}: all checks passed (spatial={spatial})')

# Gradient flow test
x = torch.randn(2, 32, 8, 8, requires_grad=True)
block = torch.nn.Sequential(
    SAPEOverwrite(2),
    Residual(torch.nn.Sequential(ChannelAffine(2, 32), AttentionND(2, 32, 4)))
)
loss = block(x).sum()
loss.backward()
assert x.grad is not None and x.grad.abs().sum() > 0, 'No gradient flow'
print('Gradient flow: OK')
