import torch
from safetensors.torch import save_file

# XNOR3 = XOR(XNOR(a,b), c)
# Using standard (non-optimized) weights

# XNOR(a,b): NOR + AND -> OR structure
# n1 (NOR): fires when both inputs 0
# n2 (AND): fires when both inputs 1
# out (OR): fires when either n1 or n2

# XOR(x,y): OR + NAND -> AND structure
# n1 (OR): fires when at least one input 1
# n2 (NAND): fires when not both inputs 1
# out (AND): fires when both n1 and n2

weights = {
    # First XNOR: a XNOR b
    'xnor1.layer1.n1.weight': torch.tensor([-1.0, -1.0]),  # NOR
    'xnor1.layer1.n1.bias': torch.tensor([0.0]),
    'xnor1.layer1.n2.weight': torch.tensor([1.0, 1.0]),    # AND
    'xnor1.layer1.n2.bias': torch.tensor([-2.0]),
    'xnor1.layer2.weight': torch.tensor([1.0, 1.0]),       # OR
    'xnor1.layer2.bias': torch.tensor([-1.0]),

    # Second XOR: (a XNOR b) XOR c
    'xor2.layer1.n1.weight': torch.tensor([1.0, 1.0]),     # OR
    'xor2.layer1.n1.bias': torch.tensor([-1.0]),
    'xor2.layer1.n2.weight': torch.tensor([-1.0, -1.0]),   # NAND
    'xor2.layer1.n2.bias': torch.tensor([1.0]),
    'xor2.layer2.weight': torch.tensor([1.0, 1.0]),        # AND
    'xor2.layer2.bias': torch.tensor([-2.0]),
}

save_file(weights, 'model.safetensors')

mag = sum(t.abs().sum().item() for t in weights.values())
print(f'Created model.safetensors')
print(f'  magnitude: {mag:.0f}')
print(f'  parameters: {sum(t.numel() for t in weights.values())}')