import torch
from safetensors.torch import save_file

weights = {}

# 4-input Fixed Priority Arbiter
# Priority: REQ0 > REQ1 > REQ2 > REQ3

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)

# Input: REQ3, REQ2, REQ1, REQ0
# Grant 0: REQ0
add_neuron('g0', [0.0, 0.0, 0.0, 1.0], -1.0)

# Grant 1: REQ1 AND NOT REQ0
add_neuron('g1', [0.0, 0.0, 1.0, -1.0], 0.0)

# Grant 2: REQ2 AND NOT REQ1 AND NOT REQ0
add_neuron('g2', [0.0, 1.0, -1.0, -1.0], 1.0)

# Grant 3: REQ3 AND NOT REQ2 AND NOT REQ1 AND NOT REQ0
add_neuron('g3', [1.0, -1.0, -1.0, -1.0], 2.0)

save_file(weights, 'model.safetensors')

def priority_arb(r3, r2, r1, r0):
    if r0:
        return 0, 0, 0, 1
    elif r1:
        return 0, 0, 1, 0
    elif r2:
        return 0, 1, 0, 0
    elif r3:
        return 1, 0, 0, 0
    return 0, 0, 0, 0

print("Verifying priority arbiter...")
errors = 0
for reqs in range(16):
    r3, r2, r1, r0 = (reqs>>3)&1, (reqs>>2)&1, (reqs>>1)&1, reqs&1
    result = priority_arb(r3, r2, r1, r0)
    grant_count = sum(result)
    if reqs > 0 and grant_count != 1:
        errors += 1
    if reqs == 0 and grant_count != 0:
        errors += 1

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}")
print(f"Parameters: {sum(t.numel() for t in weights.values())}")