import torch
from safetensors.torch import load_file

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

def overflow_detect(a_sign, b_sign, sum_sign, weights):
    """Detect signed addition overflow. Returns 1 if overflow occurred."""
    inp = torch.tensor([float(a_sign), float(b_sign), float(sum_sign)])
    n1 = int((inp @ weights['layer1.n1.weight'].T + weights['layer1.n1.bias'] >= 0).item())
    n2 = int((inp @ weights['layer1.n2.weight'].T + weights['layer1.n2.bias'] >= 0).item())
    hidden = torch.tensor([float(n1), float(n2)])
    return int((hidden @ weights['layer2.weight'].T + weights['layer2.bias'] >= 0).item())

if __name__ == '__main__':
    w = load_model()
    print('Overflow detection truth table:')
    print('a_sign b_sign sum_sign | overflow | meaning')
    print('-' * 55)
    for a in [0, 1]:
        for b in [0, 1]:
            for s in [0, 1]:
                result = overflow_detect(a, b, s, w)
                a_str = 'pos' if a == 0 else 'neg'
                b_str = 'pos' if b == 0 else 'neg'
                s_str = 'pos' if s == 0 else 'neg'
                marker = 'OVERFLOW!' if result else 'ok'
                print(f'   {a}       {b}        {s}     |    {result}     | {a_str}+{b_str}={s_str} {marker}')