import torch
from safetensors.torch import load_file

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

def full_adder(a, b, cin, w, prefix):
    """Single full adder: a + b + cin = (sum, cout)"""
    inp = torch.tensor([float(a), float(b)])

    # HA1: a + b
    or_out = (inp * w[f'{prefix}.ha1.sum.layer1.or.weight']).sum() + w[f'{prefix}.ha1.sum.layer1.or.bias'] >= 0
    nand_out = (inp * w[f'{prefix}.ha1.sum.layer1.nand.weight']).sum() + w[f'{prefix}.ha1.sum.layer1.nand.bias'] >= 0
    l1 = torch.tensor([float(or_out), float(nand_out)])
    s1 = float((l1 * w[f'{prefix}.ha1.sum.layer2.weight']).sum() + w[f'{prefix}.ha1.sum.layer2.bias'] >= 0)
    c1 = float((inp * w[f'{prefix}.ha1.carry.weight']).sum() + w[f'{prefix}.ha1.carry.bias'] >= 0)

    # HA2: s1 + cin
    inp2 = torch.tensor([s1, float(cin)])
    or_out2 = (inp2 * w[f'{prefix}.ha2.sum.layer1.or.weight']).sum() + w[f'{prefix}.ha2.sum.layer1.or.bias'] >= 0
    nand_out2 = (inp2 * w[f'{prefix}.ha2.sum.layer1.nand.weight']).sum() + w[f'{prefix}.ha2.sum.layer1.nand.bias'] >= 0
    l2 = torch.tensor([float(or_out2), float(nand_out2)])
    s = int((l2 * w[f'{prefix}.ha2.sum.layer2.weight']).sum() + w[f'{prefix}.ha2.sum.layer2.bias'] >= 0)
    c2 = float((inp2 * w[f'{prefix}.ha2.carry.weight']).sum() + w[f'{prefix}.ha2.carry.bias'] >= 0)

    # Carry OR
    cout = int((torch.tensor([c1, c2]) * w[f'{prefix}.carry_or.weight']).sum() + w[f'{prefix}.carry_or.bias'] >= 0)
    return s, cout

def ripple_carry_4bit(a, b, cin, weights):
    """4-bit ripple carry adder.
    a, b: lists of 4 bits each (LSB first)
    cin: carry in
    Returns: (sums, cout) where sums is 4-bit list (LSB first)
    """
    carries = [cin]
    sums = []
    for i in range(4):
        s, c = full_adder(a[i], b[i], carries[i], weights, f'fa{i}')
        sums.append(s)
        carries.append(c)
    return sums, carries[4]

if __name__ == '__main__':
    w = load_model()
    print('4-bit Ripple Carry Adder')
    print('Testing all combinations would be 2^9 = 512 cases.')
    print('Sample tests:')
    tests = [
        ([0,0,0,0], [0,0,0,0], 0, 0),
        ([1,0,0,0], [1,0,0,0], 0, 2),
        ([1,1,1,1], [1,0,0,0], 0, 16),
        ([1,1,1,1], [1,1,1,1], 0, 30),
        ([0,1,0,1], [0,1,0,1], 0, 20),
        ([1,1,0,0], [1,1,0,0], 0, 6),
    ]
    for a, b, cin, expected in tests:
        sums, cout = ripple_carry_4bit(a, b, cin, w)
        result = sum(bit << i for i, bit in enumerate(sums)) + (cout << 4)
        a_val = sum(bit << i for i, bit in enumerate(a))
        b_val = sum(bit << i for i, bit in enumerate(b))
        status = "OK" if result == expected else "FAIL"
        print(f'{a_val:2d} + {b_val:2d} + {cin} = {result:2d} {status}')