import torch
from safetensors.torch import save_file

weights = {}

# 4-bit Binary to BCD (for 0-9, identity; 10-15 invalid)
# Extended: detect if value >= 10 for two-digit output

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 bits
for i in range(4):
    w = [0.0] * 4
    w[i] = 1.0
    add_neuron(f'd{3-i}', w, -1.0)

# Detect >= 10 (for overflow into tens digit)
add_neuron('ge10', [8.0, 4.0, 2.0, 1.0], -10.0)

save_file(weights, 'model.safetensors')

def binary2bcd(b3, b2, b1, b0):
    val = b3*8 + b2*4 + b1*2 + b0
    if val < 10:
        return 0, b3, b2, b1, b0
    else:
        ones = val - 10
        return 1, (ones>>3)&1, (ones>>2)&1, (ones>>1)&1, ones&1

print("Verifying Binary to BCD...")
errors = 0
for b in range(16):
    b3, b2, b1, b0 = (b>>3)&1, (b>>2)&1, (b>>1)&1, b&1
    tens, d3, d2, d1, d0 = binary2bcd(b3, b2, b1, b0)
    result = tens * 10 + d3*8 + d2*4 + d1*2 + d0
    if result != b:
        errors += 1

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