import torch
from safetensors.torch import save_file

# Priority encoder: outputs binary index of highest-set input
# Inputs: I7, I6, I5, I4, I3, I2, I1, I0 (I7 is highest priority)
# Outputs: Y2, Y1, Y0 (binary encoding)

weights = {}

# Y2: fires when position >= 4 (I7 OR I6 OR I5 OR I4)
weights['y2.weight'] = torch.tensor([[1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['y2.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# Y1: fires when position has bit 1 set (positions 2,3,6,7)
# I7, I6 dominate; I5, I4 suppress; I3, I2 activate when higher bits absent
weights['y1.weight'] = torch.tensor([[16.0, 16.0, -4.0, -4.0, 1.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
weights['y1.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# Y0: fires when position has bit 0 set (positions 1,3,5,7)
# Geometric weights with alternating signs for priority cascade
weights['y0.weight'] = torch.tensor([[128.0, -64.0, 32.0, -16.0, 8.0, -4.0, 2.0, 0.0]], dtype=torch.float32)
weights['y0.bias'] = torch.tensor([-1.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

def encode8to3(i7, i6, i5, i4, i3, i2, i1, i0):
    inp = torch.tensor([float(i7), float(i6), float(i5), float(i4),
                        float(i3), float(i2), float(i1), float(i0)])
    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 y2, y1, y0

print("Verifying 8to3encoder...")
errors = 0
for val in range(256):
    bits = [(val >> (7-i)) & 1 for i in range(8)]
    y2, y1, y0 = encode8to3(*bits)

    # Expected: binary of highest set bit position
    highest = -1
    for i in range(8):
        if bits[i]:
            highest = 7 - i
            break

    if highest < 0:
        expected = (0, 0, 0)
    else:
        expected = ((highest >> 2) & 1, (highest >> 1) & 1, highest & 1)

    if (y2, y1, y0) != expected:
        errors += 1
        if errors <= 10:
            print(f"ERROR: I={''.join(map(str,bits))} -> ({y2},{y1},{y0}), expected {expected} (pos {highest})")

if errors == 0:
    print("All 256 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}")