"""Load model.safetensors and run the gate network (self-contained)."""

import torch
from safetensors.torch import load_file

GATES = ("AND", "OR", "XOR", "NAND", "NOR", "XNOR", "NOT", "BUF", "CNOT", "SWAP")
W = load_file("model.safetensors")


def forward(gate: str, a: int, b: int) -> torch.Tensor:
    x = torch.zeros(12)
    x[GATES.index(gate)] = 1.0
    x[-2], x[-1] = a, b
    for i in (0, 2, 4):
        x = torch.tanh(x @ W[f"net.{i}.weight"].T + W[f"net.{i}.bias"])
    return torch.sigmoid(x @ W["net.6.weight"].T + W["net.6.bias"])


if __name__ == "__main__":
    for gate in GATES:
        p1, p2 = forward(gate, 1, 0).tolist()
        print(f"{gate:4}  a=1 b=0  ->  ({int(p1 >= 0.5)}, {int(p2 >= 0.5)})")