import torch
from safetensors.torch import save_file

# Tree structure: parity(a,b,c,d) = XOR(XOR(a,b), XOR(c,d))
# Each XOR uses OR-NAND-AND structure

def xor_block(prefix):
    return {
        f'{prefix}.or.weight': torch.tensor([1.0, 1.0], dtype=torch.float32),
        f'{prefix}.or.bias': torch.tensor([-1.0], dtype=torch.float32),
        f'{prefix}.nand.weight': torch.tensor([-1.0, -1.0], dtype=torch.float32),
        f'{prefix}.nand.bias': torch.tensor([1.0], dtype=torch.float32),
        f'{prefix}.and.weight': torch.tensor([1.0, 1.0], dtype=torch.float32),
        f'{prefix}.and.bias': torch.tensor([-2.0], dtype=torch.float32),
    }

weights = {}
weights.update(xor_block('xor_ab'))   # XOR(a, b)
weights.update(xor_block('xor_cd'))   # XOR(c, d)
weights.update(xor_block('xor_final')) # XOR(xor_ab, xor_cd)

save_file(weights, 'model.safetensors')

def xor2(a, b, prefix):
    or_out = int(a * weights[f'{prefix}.or.weight'][0] + b * weights[f'{prefix}.or.weight'][1] + weights[f'{prefix}.or.bias'] >= 0)
    nand_out = int(a * weights[f'{prefix}.nand.weight'][0] + b * weights[f'{prefix}.nand.weight'][1] + weights[f'{prefix}.nand.bias'] >= 0)
    return int(or_out * weights[f'{prefix}.and.weight'][0] + nand_out * weights[f'{prefix}.and.weight'][1] + weights[f'{prefix}.and.bias'] >= 0)

def parity4(a, b, c, d):
    xor_ab = xor2(a, b, 'xor_ab')
    xor_cd = xor2(c, d, 'xor_cd')
    return xor2(xor_ab, xor_cd, 'xor_final')

print("Verifying parity4...")
errors = 0
for i in range(16):
    a, b, c, d = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
    result = parity4(a, b, c, d)
    expected = a ^ b ^ c ^ d
    if result != expected:
        errors += 1
        print(f"ERROR: parity({a},{b},{c},{d}) = {result}, expected {expected}")
if errors == 0:
    print("All 16 test cases passed!")
print(f"Magnitude: {sum(t.abs().sum().item() for t in weights.values()):.0f}")