import torch
from safetensors.torch import save_file

weights = {}

# Weighted Threshold Function
# Inputs: x3, x2, x1, x0 with weights 4, 3, 2, 1
# Output: y = 1 iff 4*x3 + 3*x2 + 2*x1 + 1*x0 >= 6
#
# This represents weighted voting where different inputs have different influence.
# Maximum sum = 4+3+2+1 = 10
# Threshold = 6 (need majority of weighted votes)
#
# Single neuron implementation!

weights['y.weight'] = torch.tensor([[4.0, 3.0, 2.0, 1.0]], dtype=torch.float32)
weights['y.bias'] = torch.tensor([-6.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

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

def reference(x3, x2, x1, x0):
    weighted_sum = 4*x3 + 3*x2 + 2*x1 + 1*x0
    return 1 if weighted_sum >= 6 else 0

print("Verifying Weighted Threshold (4,3,2,1 >= 6)...")
errors = 0
for i in range(16):
    x3, x2, x1, x0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
    result = weighted_threshold(x3, x2, x1, x0)
    expected = reference(x3, x2, x1, x0)
    if result != expected:
        errors += 1
        print(f"ERROR: ({x3},{x2},{x1},{x0}) -> {result}, expected {expected}")

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

print("\nTruth Table:")
print("x3 x2 x1 x0 | w_sum | y")
print("-" * 26)
for i in range(16):
    x3, x2, x1, x0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
    y = weighted_threshold(x3, x2, x1, x0)
    ws = 4*x3 + 3*x2 + 2*x1 + 1*x0
    print(f" {x3}  {x2}  {x1}  {x0}  |   {ws:2d}  | {y}")

mag = sum(t.abs().sum().item() for t in weights.values())
print(f"\nMagnitude: {mag:.0f}")
print(f"Parameters: {sum(t.numel() for t in weights.values())}")
print(f"Neurons: {len([k for k in weights.keys() if 'weight' in k])}")