import torch
from safetensors.torch import save_file

weights = {}

# Each output yi = NOT(ai)
# NOT(x): weight -1, bias 0, fires when x = 0
for i in range(4):
    w = [0.0, 0.0, 0.0, 0.0]
    w[i] = -1.0
    weights[f'y{i}.weight'] = torch.tensor([w], dtype=torch.float32)
    weights[f'y{i}.bias'] = torch.tensor([0.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

# Verify
def negator4(a3, a2, a1, a0):
    inp = torch.tensor([float(a3), float(a2), float(a1), float(a0)])
    outputs = []
    for i in range(4):
        y = int((inp * weights[f'y{i}.weight']).sum() + weights[f'y{i}.bias'] >= 0)
        outputs.append(y)
    return outputs

print("Verifying negator4bit...")
errors = 0
for i in range(16):
    a3, a2, a1, a0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
    result = negator4(a3, a2, a1, a0)
    expected = [1 - a3, 1 - a2, 1 - a1, 1 - a0]
    if result != expected:
        errors += 1
        print(f"ERROR: {a3}{a2}{a1}{a0} -> {result}, expected {expected}")

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