import torch
from fractal_block import TwoChannelFractalBlock

device = 'cuda' if torch.cuda.is_available() else 'cpu'
try:
    m = TwoChannelFractalBlock(8).to(device)
    a = torch.randn(4, 3, 8, device=device)
    b = torch.randn(4, 3, 8, device=device)
    ap, bp = m(a, b)
    assert ap.shape == bp.shape == (4, 8)
    assert torch.isfinite(ap).all() and torch.isfinite(bp).all()
    print({'device': device, 'a_parent': tuple(ap.shape), 'b_parent': tuple(bp.shape), 'coefficients': m.coefficients()})
except Exception as e:
    if device == 'cuda':
        m = TwoChannelFractalBlock(8)
        ap, bp = m(torch.randn(4,3,8), torch.randn(4,3,8))
        print({'device': 'cpu-fallback', 'a_parent': tuple(ap.shape), 'b_parent': tuple(bp.shape), 'error': str(e)})
    else:
        raise