import torch
from safetensors.torch import save_file

weights = {}

# y3 = x3 (at least 1 of [x3])
weights['y3.weight'] = torch.tensor([[1.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['y3.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# y2 = x3 OR x2 (at least 1 of [x3,x2])
weights['y2.weight'] = torch.tensor([[1.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
weights['y2.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# y1 = x3 OR x2 OR x1
weights['y1.weight'] = torch.tensor([[1.0, 1.0, 1.0, 0.0]], dtype=torch.float32)
weights['y1.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# y0 = x3 OR x2 OR x1 OR x0
weights['y0.weight'] = torch.tensor([[1.0, 1.0, 1.0, 1.0]], dtype=torch.float32)
weights['y0.bias'] = torch.tensor([-1.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

def prefix_or(x3, x2, x1, x0):
    inp = torch.tensor([float(x3), float(x2), float(x1), float(x0)])
    y3 = int((inp @ weights['y3.weight'].T + weights['y3.bias'] >= 0).item())
    y2 = int((inp @ weights['y2.weight'].T + weights['y2.bias'] >= 0).item())
    y1 = int((inp @ weights['y1.weight'].T + weights['y1.bias'] >= 0).item())
    y0 = int((inp @ weights['y0.weight'].T + weights['y0.bias'] >= 0).item())
    return y3, y2, y1, y0

print("Verifying prefix-or...")
errors = 0
for i in range(16):
    x3, x2, x1, x0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
    y3, y2, y1, y0 = prefix_or(x3, x2, x1, x0)
    exp_y3 = x3
    exp_y2 = x3 | x2
    exp_y1 = x3 | x2 | x1
    exp_y0 = x3 | x2 | x1 | x0
    if (y3, y2, y1, y0) != (exp_y3, exp_y2, exp_y1, exp_y0):
        errors += 1
if errors == 0:
    print("All 16 test cases passed!")
else:
    print(f"FAILED: {errors} errors")
print(f"Magnitude: {sum(t.abs().sum().item() for t in weights.values()):.0f}")