import torch
from safetensors.torch import save_file

weights = {}

# Single digit BCD (4 bits) to Binary (4 bits)
# BCD: 0-9 encoded in 4 bits
# Output: same value (identity for single digit)

def add_neuron(name, w_list, bias):
    weights[f'{name}.weight'] = torch.tensor([w_list], dtype=torch.float32)
    weights[f'{name}.bias'] = torch.tensor([bias], dtype=torch.float32)

# Pass through (BCD 0-9 is same as binary 0-9)
for i in range(4):
    w = [0.0] * 4
    w[i] = 1.0
    add_neuron(f'b{3-i}', w, -1.0)

save_file(weights, 'model.safetensors')

def bcd2binary(d3, d2, d1, d0):
    return d3, d2, d1, d0

print("Verifying BCD to Binary...")
errors = 0
for d in range(10):  # Valid BCD: 0-9
    d3, d2, d1, d0 = (d>>3)&1, (d>>2)&1, (d>>1)&1, d&1
    b3, b2, b1, b0 = bcd2binary(d3, d2, d1, d0)
    result = b3*8 + b2*4 + b1*2 + b0
    if result != d:
        errors += 1

if errors == 0:
    print("All 10 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())}")