import torch
from safetensors.torch import save_file

weights = {}

# 4-bit Han-Carlson Parallel Prefix Adder
# Hybrid of Kogge-Stone and Brent-Kung

def add_and(name, idx_a, idx_b, n_inputs):
    w = [0.0] * n_inputs
    w[idx_a] = 1.0
    w[idx_b] = 1.0
    weights[f'{name}.weight'] = torch.tensor([w], dtype=torch.float32)
    weights[f'{name}.bias'] = torch.tensor([-2.0], dtype=torch.float32)

def add_xor_stage1(name, idx_a, idx_b, n_inputs):
    w_or = [0.0] * n_inputs
    w_or[idx_a] = 1.0
    w_or[idx_b] = 1.0
    weights[f'{name}.or.weight'] = torch.tensor([w_or], dtype=torch.float32)
    weights[f'{name}.or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
    w_nand = [0.0] * n_inputs
    w_nand[idx_a] = -1.0
    w_nand[idx_b] = -1.0
    weights[f'{name}.nand.weight'] = torch.tensor([w_nand], dtype=torch.float32)
    weights[f'{name}.nand.bias'] = torch.tensor([1.0], dtype=torch.float32)

def add_xor_stage2(name):
    weights[f'{name}.and.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
    weights[f'{name}.and.bias'] = torch.tensor([-2.0], dtype=torch.float32)

for i in range(4):
    a_idx = 3 - i
    b_idx = 7 - i
    add_and(f'g{i}', a_idx, b_idx, 9)
    add_xor_stage1(f'p{i}', a_idx, b_idx, 9)

for i in range(4):
    add_xor_stage2(f'p{i}')

save_file(weights, 'model.safetensors')

def han_carlson_add(a3, a2, a1, a0, b3, b2, b1, b0, cin):
    a = [a0, a1, a2, a3]
    b = [b0, b1, b2, b3]

    g = [a[i] & b[i] for i in range(4)]
    p = [a[i] ^ b[i] for i in range(4)]

    # Level 1: odd positions
    g10 = g[1] | (p[1] & g[0])
    p10 = p[1] & p[0]
    g32 = g[3] | (p[3] & g[2])
    p32 = p[3] & p[2]

    # Level 2: Kogge-Stone style for odd
    g30 = g32 | (p32 & g10)

    # Level 3: back-propagate
    g20 = g[2] | (p[2] & g10)

    c0 = g[0] | (p[0] & cin)
    c1 = g10 | (p10 & cin)
    c2 = g20 | (p[2] & p10 & cin)
    c3 = g30 | (p32 & p10 & cin)

    s0 = p[0] ^ cin
    s1 = p[1] ^ c0
    s2 = p[2] ^ c1
    s3 = p[3] ^ c2

    return s3, s2, s1, s0, c3

print("Verifying 4-bit Han-Carlson adder...")
errors = 0
for a in range(16):
    for b in range(16):
        for cin in range(2):
            a3, a2, a1, a0 = (a>>3)&1, (a>>2)&1, (a>>1)&1, a&1
            b3, b2, b1, b0 = (b>>3)&1, (b>>2)&1, (b>>1)&1, b&1
            s3, s2, s1, s0, cout = han_carlson_add(a3, a2, a1, a0, b3, b2, b1, b0, cin)
            result = cout*16 + s3*8 + s2*4 + s1*2 + s0
            expected = a + b + cin
            if result != expected:
                errors += 1
                if errors <= 3:
                    print(f"ERROR: {a}+{b}+{cin} = {result}, expected {expected}")

if errors == 0:
    print("All 512 test cases passed!")
else:
    print(f"FAILED: {errors} errors")

mag = sum(t.abs().sum().item() for t in weights.values())
print(f"Magnitude: {mag:.0f}")
print(f"Parameters: {sum(t.numel() for t in weights.values())}")