import torch
from safetensors.torch import load_file

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

def decrementer4(a3, a2, a1, a0, weights):
    """4-bit decrementer: returns (input - 1) mod 16"""
    inp = torch.tensor([float(a3), float(a2), float(a1), float(a0)])

    # Layer 1
    y0 = int((inp @ weights['y0.weight'].T + weights['y0.bias'] >= 0).item())
    b1 = int((inp @ weights['b1.weight'].T + weights['b1.bias'] >= 0).item())
    b2 = int((inp @ weights['b2.weight'].T + weights['b2.bias'] >= 0).item())
    y1_and = int((inp @ weights['y1_and.weight'].T + weights['y1_and.bias'] >= 0).item())

    # Layer 2
    l2_in = torch.tensor([float(a3), float(a2), float(b1), float(b2), float(y1_and)])
    y1 = int((l2_in @ weights['y1.weight'].T + weights['y1.bias'] >= 0).item())
    y2_or = int((l2_in @ weights['y2_or.weight'].T + weights['y2_or.bias'] >= 0).item())
    y2_nand = int((l2_in @ weights['y2_nand.weight'].T + weights['y2_nand.bias'] >= 0).item())
    y3_or = int((l2_in @ weights['y3_or.weight'].T + weights['y3_or.bias'] >= 0).item())
    y3_nand = int((l2_in @ weights['y3_nand.weight'].T + weights['y3_nand.bias'] >= 0).item())

    # Layer 3
    l3_y2 = torch.tensor([float(y2_or), float(y2_nand)])
    l3_y3 = torch.tensor([float(y3_or), float(y3_nand)])
    y2 = int((l3_y2 @ weights['y2.weight'].T + weights['y2.bias'] >= 0).item())
    y3 = int((l3_y3 @ weights['y3.weight'].T + weights['y3.bias'] >= 0).item())

    return [y3, y2, y1, y0]

if __name__ == '__main__':
    w = load_model()
    print('Decrementer4bit:')
    for i in range(16):
        a3, a2, a1, a0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
        result = decrementer4(a3, a2, a1, a0, w)
        out_val = result[0]*8 + result[1]*4 + result[2]*2 + result[3]
        print(f'  {i:2d} ({a3}{a2}{a1}{a0}) - 1 = {out_val:2d} ({"".join(map(str, result))})')