import torch
from safetensors.torch import load_file

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

def clz8(a7, a6, a5, a4, a3, a2, a1, a0, weights):
    """8-bit count leading zeros. Returns 4-bit binary encoding of 0-8."""
    inp = torch.tensor([float(a7), float(a6), float(a5), float(a4),
                        float(a3), float(a2), float(a1), float(a0)])

    # Layer 1: priority detection
    has7 = int((inp @ weights['has7.weight'].T + weights['has7.bias'] >= 0).item())
    has6_first = int((inp @ weights['has6_first.weight'].T + weights['has6_first.bias'] >= 0).item())
    has5_first = int((inp @ weights['has5_first.weight'].T + weights['has5_first.bias'] >= 0).item())
    has4_first = int((inp @ weights['has4_first.weight'].T + weights['has4_first.bias'] >= 0).item())
    has3_first = int((inp @ weights['has3_first.weight'].T + weights['has3_first.bias'] >= 0).item())
    has2_first = int((inp @ weights['has2_first.weight'].T + weights['has2_first.bias'] >= 0).item())
    has1_first = int((inp @ weights['has1_first.weight'].T + weights['has1_first.bias'] >= 0).item())
    has0_first = int((inp @ weights['has0_first.weight'].T + weights['has0_first.bias'] >= 0).item())
    all_zero = int((inp @ weights['all_zero.weight'].T + weights['all_zero.bias'] >= 0).item())

    # Layer 2: binary encoding
    l1 = torch.tensor([float(has7), float(has6_first), float(has5_first), float(has4_first),
                       float(has3_first), float(has2_first), float(has1_first), float(has0_first),
                       float(all_zero)])
    y0 = int((l1 @ weights['y0.weight'].T + weights['y0.bias'] >= 0).item())
    y1 = int((l1 @ weights['y1.weight'].T + weights['y1.bias'] >= 0).item())
    y2 = int((l1 @ weights['y2.weight'].T + weights['y2.bias'] >= 0).item())
    y3 = int((l1 @ weights['y3.weight'].T + weights['y3.bias'] >= 0).item())

    return [y3, y2, y1, y0]

if __name__ == '__main__':
    w = load_model()
    print('clz8 examples:')
    test_cases = [
        (1, 0, 0, 0, 0, 0, 0, 0),  # 0 leading zeros
        (0, 1, 0, 0, 0, 0, 0, 0),  # 1 leading zero
        (0, 0, 0, 0, 1, 0, 0, 0),  # 4 leading zeros
        (0, 0, 0, 0, 0, 0, 0, 1),  # 7 leading zeros
        (0, 0, 0, 0, 0, 0, 0, 0),  # 8 leading zeros
    ]
    for bits in test_cases:
        result = clz8(*bits, w)
        clz_val = result[0]*8 + result[1]*4 + result[2]*2 + result[3]
        print(f'  {bits} -> {result} = {clz_val}')