import torch
from safetensors.torch import load_file

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

def t_flipflop(t, q_prev, weights):
    """T Flip-Flop: T=1 toggles, T=0 holds."""
    inp = torch.tensor([float(t), float(q_prev)])

    or_out = int((inp @ weights['or.weight'].T + weights['or.bias'] >= 0).item())
    nand_out = int((inp @ weights['nand.weight'].T + weights['nand.bias'] >= 0).item())
    nor_out = int((inp @ weights['nor.weight'].T + weights['nor.bias'] >= 0).item())
    and_out = int((inp @ weights['and.weight'].T + weights['and.bias'] >= 0).item())

    l1_q = torch.tensor([float(or_out), float(nand_out)])
    q = int((l1_q @ weights['q.weight'].T + weights['q.bias'] >= 0).item())

    l1_qn = torch.tensor([float(nor_out), float(and_out)])
    qn = int((l1_qn @ weights['qn.weight'].T + weights['qn.bias'] >= 0).item())

    return q, qn

if __name__ == '__main__':
    w = load_model()
    print('T Flip-Flop (counter demo):')
    q = 0
    for i in range(8):
        print(f'  Step {i}: Q={q}')
        q, _ = t_flipflop(1, q, w)  # Always toggle