import torch
from safetensors.torch import load_file

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

def signextend4to8(x3, x2, x1, x0, weights):
    """Sign extend 4-bit signed integer to 8-bit.
    Input: x3 (MSB/sign), x2, x1, x0
    Output: [y0, y1, ..., y7] where y7 is MSB
    """
    inp = torch.tensor([float(x3), float(x2), float(x1), float(x0)])
    outputs = []
    for i in range(8):
        y = int((inp * weights[f'y{i}.weight']).sum() + weights[f'y{i}.bias'] >= 0)
        outputs.append(y)
    return outputs

if __name__ == '__main__':
    w = load_model()
    print('signextend4to8 truth table:')
    print('4-bit (dec) | 8-bit binary    | 8-bit (dec)')
    print('------------+-----------------+------------')
    for val in range(16):
        x3, x2, x1, x0 = (val >> 3) & 1, (val >> 2) & 1, (val >> 1) & 1, val & 1
        result = signextend4to8(x3, x2, x1, x0, w)
        signed_4 = val if val < 8 else val - 16
        result_val = sum(b << i for i, b in enumerate(result))
        signed_8 = result_val if result_val < 128 else result_val - 256
        bits = ''.join(str(b) for b in reversed(result))
        print(f'  {signed_4:3d}       |    {bits}    |   {signed_8:4d}')