import torch
from safetensors.torch import save_file

# Gray to Binary: B[i] = XOR(G[n-1], G[n-2], ..., G[i])
# B3 = G3
# B2 = G3 XOR G2
# B1 = G3 XOR G2 XOR G1
# B0 = G3 XOR G2 XOR G1 XOR G0
#
# Using cascade XOR: XOR(a,b) = AND(OR(a,b), NAND(a,b))
# With sharing: X1 = XOR(G3,G2), X2 = XOR(X1,G1), X3 = XOR(X2,G0)
# B3=G3, B2=X1, B1=X2, B0=X3

weights = {}

# Inputs: G3, G2, G1, G0

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

# === X1 = XOR(G3, G2) → B2 ===
# Layer 1: OR and NAND from inputs
weights['x1_or.weight'] = torch.tensor([[1.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
weights['x1_or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
weights['x1_nand.weight'] = torch.tensor([[-1.0, -1.0, 0.0, 0.0]], dtype=torch.float32)
weights['x1_nand.bias'] = torch.tensor([1.0], dtype=torch.float32)
# Layer 2: AND(or, nand) → X1 = B2
weights['b2.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['b2.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# === X2 = XOR(X1, G1) → B1 ===
# Need to feed X1 (from layer 2) and G1 (from input)
# Layer 3: OR(X1, G1), NAND(X1, G1)
# X1 is neuron output, G1 is input[2]
weights['x2_or.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)  # [X1, G1]
weights['x2_or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
weights['x2_nand.weight'] = torch.tensor([[-1.0, -1.0]], dtype=torch.float32)
weights['x2_nand.bias'] = torch.tensor([1.0], dtype=torch.float32)
# Layer 4: AND → X2 = B1
weights['b1.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['b1.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# === X3 = XOR(X2, G0) → B0 ===
# Layer 5: OR(X2, G0), NAND(X2, G0)
weights['x3_or.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['x3_or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
weights['x3_nand.weight'] = torch.tensor([[-1.0, -1.0]], dtype=torch.float32)
weights['x3_nand.bias'] = torch.tensor([1.0], dtype=torch.float32)
# Layer 6: AND → X3 = B0
weights['b0.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['b0.bias'] = torch.tensor([-2.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

def gray2binary(g3, g2, g1, g0):
    # B3 = G3
    b3 = int(g3 * 2 - 1 >= 0)

    # X1 = XOR(G3, G2)
    x1_or = int(g3 + g2 - 1 >= 0)
    x1_nand = int(-g3 - g2 + 1 >= 0)
    x1 = int(x1_or + x1_nand - 2 >= 0)
    b2 = x1

    # X2 = XOR(X1, G1)
    x2_or = int(x1 + g1 - 1 >= 0)
    x2_nand = int(-x1 - g1 + 1 >= 0)
    x2 = int(x2_or + x2_nand - 2 >= 0)
    b1 = x2

    # X3 = XOR(X2, G0)
    x3_or = int(x2 + g0 - 1 >= 0)
    x3_nand = int(-x2 - g0 + 1 >= 0)
    x3 = int(x3_or + x3_nand - 2 >= 0)
    b0 = x3

    return b3, b2, b1, b0

print("Verifying gray2binary...")
errors = 0
for i in range(16):
    # Convert i to gray code
    gray = i ^ (i >> 1)
    g3, g2, g1, g0 = (gray >> 3) & 1, (gray >> 2) & 1, (gray >> 1) & 1, gray & 1

    b3, b2, b1, b0 = gray2binary(g3, g2, g1, g0)
    result = b3 * 8 + b2 * 4 + b1 * 2 + b0

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

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}")