import torch
from safetensors.torch import save_file

# Binary to Gray: G[i] = B[i] XOR B[i+1] (where B[n]=0)
# G3 = B3
# G2 = B3 XOR B2
# G1 = B2 XOR B1
# G0 = B1 XOR B0
#
# Each XOR uses 3 neurons: OR, NAND, AND

weights = {}

# Inputs: B3, B2, B1, B0

# === G3 = B3 (identity) ===
weights['g3.weight'] = torch.tensor([[2.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['g3.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# === G2 = XOR(B3, B2) ===
weights['g2_or.weight'] = torch.tensor([[1.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
weights['g2_or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
weights['g2_nand.weight'] = torch.tensor([[-1.0, -1.0, 0.0, 0.0]], dtype=torch.float32)
weights['g2_nand.bias'] = torch.tensor([1.0], dtype=torch.float32)
weights['g2.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['g2.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# === G1 = XOR(B2, B1) ===
weights['g1_or.weight'] = torch.tensor([[0.0, 1.0, 1.0, 0.0]], dtype=torch.float32)
weights['g1_or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
weights['g1_nand.weight'] = torch.tensor([[0.0, -1.0, -1.0, 0.0]], dtype=torch.float32)
weights['g1_nand.bias'] = torch.tensor([1.0], dtype=torch.float32)
weights['g1.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['g1.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# === G0 = XOR(B1, B0) ===
weights['g0_or.weight'] = torch.tensor([[0.0, 0.0, 1.0, 1.0]], dtype=torch.float32)
weights['g0_or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
weights['g0_nand.weight'] = torch.tensor([[0.0, 0.0, -1.0, -1.0]], dtype=torch.float32)
weights['g0_nand.bias'] = torch.tensor([1.0], dtype=torch.float32)
weights['g0.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['g0.bias'] = torch.tensor([-2.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

def binary2gray(b3, b2, b1, b0):
    inp = [b3, b2, b1, b0]

    # G3 = B3
    g3 = int(2*b3 - 1 >= 0)

    # G2 = XOR(B3, B2)
    g2_or = int(b3 + b2 - 1 >= 0)
    g2_nand = int(-b3 - b2 + 1 >= 0)
    g2 = int(g2_or + g2_nand - 2 >= 0)

    # G1 = XOR(B2, B1)
    g1_or = int(b2 + b1 - 1 >= 0)
    g1_nand = int(-b2 - b1 + 1 >= 0)
    g1 = int(g1_or + g1_nand - 2 >= 0)

    # G0 = XOR(B1, B0)
    g0_or = int(b1 + b0 - 1 >= 0)
    g0_nand = int(-b1 - b0 + 1 >= 0)
    g0 = int(g0_or + g0_nand - 2 >= 0)

    return g3, g2, g1, g0

print("Verifying binary2gray...")
errors = 0
for i in range(16):
    b3, b2, b1, b0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
    g3, g2, g1, g0 = binary2gray(b3, b2, b1, b0)
    result = g3 * 8 + g2 * 4 + g1 * 2 + g0

    expected = i ^ (i >> 1)  # Standard gray code formula

    if result != expected:
        errors += 1
        print(f"ERROR: binary {b3}{b2}{b1}{b0} -> gray {g3}{g2}{g1}{g0} (={result}), expected {expected}")

if errors == 0:
    print("All 16 test cases passed!")

mag = sum(t.abs().sum().item() for t in weights.values())
print(f"Magnitude: {mag:.0f}")