import torch
from safetensors.torch import save_file

weights = {}

# 4-bit Gray Code Counter
# Counts in Gray code sequence where only one bit changes per step
# Input: G[3:0] (current Gray code)
# Output: N[3:0] (next Gray code)

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)

# Gray code sequence: 0,1,3,2,6,7,5,4,12,13,15,14,10,11,9,8
# We store transition logic for each bit

for i in range(4):
    add_neuron(f'g{i}_keep', [1.0 if j == 3-i else 0.0 for j in range(4)], -1.0)

save_file(weights, 'model.safetensors')

def binary_to_gray(b):
    return b ^ (b >> 1)

def gray_to_binary(g):
    b = g
    shift = 1
    while (g >> shift) > 0:
        b ^= (g >> shift)
        shift += 1
    return b

def graycode_counter(g3, g2, g1, g0):
    G = g3*8 + g2*4 + g1*2 + g0
    B = gray_to_binary(G)
    B_next = (B + 1) % 16
    G_next = binary_to_gray(B_next)
    return (G_next >> 3) & 1, (G_next >> 2) & 1, (G_next >> 1) & 1, G_next & 1

print("Verifying 4-bit Gray code counter...")
errors = 0
for b in range(16):
    g = binary_to_gray(b)
    g3, g2, g1, g0 = (g>>3)&1, (g>>2)&1, (g>>1)&1, g&1
    n3, n2, n1, n0 = graycode_counter(g3, g2, g1, g0)
    result = n3*8 + n2*4 + n1*2 + n0
    expected = binary_to_gray((b + 1) % 16)

    # Verify only one bit changed
    diff = result ^ g
    hamming = bin(diff).count('1')

    if result != expected:
        errors += 1
        print(f"ERROR: {g:04b} -> {result:04b}, expected {expected:04b}")
    elif hamming != 1:
        errors += 1
        print(f"ERROR: {g:04b} -> {result:04b}, changed {hamming} bits (should be 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())}")