import torch
from safetensors.torch import load_file

def load_model(path='model.safetensors'):
    return load_file(path)

def binary2gray(b3, b2, b1, b0, weights):
    """Convert 4-bit binary to Gray code."""
    inp = torch.tensor([float(b3), float(b2), float(b1), float(b0)])

    # G3 = B3
    g3 = int((inp @ weights['g3.weight'].T + weights['g3.bias'] >= 0).item())

    # G2 = XOR(B3, B2)
    g2_or = int((inp @ weights['g2_or.weight'].T + weights['g2_or.bias'] >= 0).item())
    g2_nand = int((inp @ weights['g2_nand.weight'].T + weights['g2_nand.bias'] >= 0).item())
    g2_vec = torch.tensor([float(g2_or), float(g2_nand)])
    g2 = int((g2_vec @ weights['g2.weight'].T + weights['g2.bias'] >= 0).item())

    # G1 = XOR(B2, B1)
    g1_or = int((inp @ weights['g1_or.weight'].T + weights['g1_or.bias'] >= 0).item())
    g1_nand = int((inp @ weights['g1_nand.weight'].T + weights['g1_nand.bias'] >= 0).item())
    g1_vec = torch.tensor([float(g1_or), float(g1_nand)])
    g1 = int((g1_vec @ weights['g1.weight'].T + weights['g1.bias'] >= 0).item())

    # G0 = XOR(B1, B0)
    g0_or = int((inp @ weights['g0_or.weight'].T + weights['g0_or.bias'] >= 0).item())
    g0_nand = int((inp @ weights['g0_nand.weight'].T + weights['g0_nand.bias'] >= 0).item())
    g0_vec = torch.tensor([float(g0_or), float(g0_nand)])
    g0 = int((g0_vec @ weights['g0.weight'].T + weights['g0.bias'] >= 0).item())

    return g3, g2, g1, g0

if __name__ == '__main__':
    w = load_model()
    print('Binary to Gray conversion:')
    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, w)
        print(f'  {b3}{b2}{b1}{b0} ({i:2d}) -> {g3}{g2}{g1}{g0}')