import torch
from safetensors.torch import save_file

# XNOR4 = XNOR(XNOR(a,b), XNOR(c,d))
# Tree structure using 3 XNOR gates
# Each XNOR: NOR + AND -> OR

def make_xnor_weights(prefix):
    return {
        f'{prefix}.layer1.n1.weight': torch.tensor([-1.0, -1.0]),  # NOR
        f'{prefix}.layer1.n1.bias': torch.tensor([0.0]),
        f'{prefix}.layer1.n2.weight': torch.tensor([1.0, 1.0]),    # AND
        f'{prefix}.layer1.n2.bias': torch.tensor([-2.0]),
        f'{prefix}.layer2.weight': torch.tensor([1.0, 1.0]),       # OR
        f'{prefix}.layer2.bias': torch.tensor([-1.0]),
    }

weights = {}
weights.update(make_xnor_weights('xnor1'))  # XNOR(a,b)
weights.update(make_xnor_weights('xnor2'))  # XNOR(c,d)
weights.update(make_xnor_weights('xnor3'))  # XNOR(xnor_ab, xnor_cd)

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())}')