import torch
from safetensors.torch import save_file

weights = {}

# CRC-16-CCITT (polynomial 0x1021: x^16 + x^12 + x^5 + 1)
# 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[15:0] (16 bits), D (1 bit) = 17 inputs total

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

# Shift neurons
for i in range(16):
    w = [0.0] * 17
    if i < 15:
        w[i+1] = 1.0  # Shift from next position
    add_neuron(f'c{i}_shift', w, -1.0 if i < 15 else 0.0)

save_file(weights, 'model.safetensors')

def crc16_step(crc, data_bit, poly=0x1021):
    feedback = ((crc >> 15) ^ data_bit) & 1
    crc = (crc << 1) & 0xFFFF
    if feedback:
        crc ^= poly
    return crc

print("Verifying CRC-16 step function...")
errors = 0
# Test subset (full test would be 2^17 cases)
for crc in [0, 0x1234, 0xFFFF, 0x8000, 0x0001]:
    for d in [0, 1]:
        result = crc16_step(crc, d)
        fb = ((crc >> 15) ^ d) & 1
        expected = ((crc << 1) & 0xFFFF) ^ (0x1021 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, 1, 0, 1, 0, 1]
crc = 0xFFFF  # Standard initial value
for bit in msg:
    crc = crc16_step(crc, bit)
print(f"CRC-16 of message: 0x{crc:04X}")

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