import torch
from safetensors.torch import save_file

weights = {}

# CRC-32 (IEEE 802.3 polynomial)
# Polynomial: 0x04C11DB7
# Computes one step of CRC given current CRC state and input bit

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)

# Inputs: C[31:0] (32 bits), D (1 bit) = 33 inputs

# Feedback = C31 XOR D
add_neuron('fb_or', [1.0] + [0.0]*31 + [1.0], -1.0)
add_neuron('fb_nand', [-1.0] + [0.0]*31 + [-1.0], 1.0)

# Shift neurons
for i in range(32):
    w = [0.0] * 33
    if i < 31:
        w[i+1] = 1.0
    add_neuron(f'c{i}_shift', w, -1.0 if i < 31 else 0.0)

save_file(weights, 'model.safetensors')

def crc32_step(crc, data_bit, poly=0x04C11DB7):
    feedback = ((crc >> 31) ^ data_bit) & 1
    crc = (crc << 1) & 0xFFFFFFFF
    if feedback:
        crc ^= poly
    return crc

print("Verifying CRC-32 step function...")
errors = 0
for crc in [0, 0xFFFFFFFF, 0x12345678, 0x80000000]:
    for d in [0, 1]:
        result = crc32_step(crc, d)
        fb = ((crc >> 31) ^ d) & 1
        expected = ((crc << 1) & 0xFFFFFFFF) ^ (0x04C11DB7 if fb else 0)
        if result != expected:
            errors += 1

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

# Test with message
msg = [0, 1, 0, 0, 1, 0, 0, 1]  # ASCII 'I'
crc = 0xFFFFFFFF
for bit in msg:
    crc = crc32_step(crc, bit)
print(f"CRC-32 partial: 0x{crc:08X}")

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())}")