import torch
from safetensors.torch import save_file

weights = {}

# 4-bit Unsigned Saturating Adder
# Inputs: a3,a2,a1,a0, b3,b2,b1,b0 (8 inputs)
# Outputs: s3,s2,s1,s0, saturated (5 outputs)
#
# If A + B > 15: output 15 (1111), saturated=1
# Else: output A + B, saturated=0

def add_xor(name):
    weights[f'{name}.or.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
    weights[f'{name}.or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
    weights[f'{name}.nand.weight'] = torch.tensor([[-1.0, -1.0]], dtype=torch.float32)
    weights[f'{name}.nand.bias'] = torch.tensor([1.0], dtype=torch.float32)
    weights[f'{name}.and.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
    weights[f'{name}.and.bias'] = torch.tensor([-2.0], dtype=torch.float32)

def add_fa(name):
    add_xor(f'{name}.xor1')
    add_xor(f'{name}.sum')
    weights[f'{name}.carry.weight'] = torch.tensor([[1.0, 1.0, 1.0]], dtype=torch.float32)
    weights[f'{name}.carry.bias'] = torch.tensor([-2.0], dtype=torch.float32)

def add_mux(name):
    weights[f'{name}.sel0.weight'] = torch.tensor([[1.0, -1.0]], dtype=torch.float32)
    weights[f'{name}.sel0.bias'] = torch.tensor([-2.0], dtype=torch.float32)
    weights[f'{name}.sel1.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
    weights[f'{name}.sel1.bias'] = torch.tensor([-2.0], dtype=torch.float32)
    weights[f'{name}.or.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
    weights[f'{name}.or.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# 4-bit ripple-carry adder
for i in range(4):
    add_fa(f'fa{i}')

# Saturation logic: if carry out, output all 1s
# Each output bit: MUX(sum_bit, 1, cout)
for i in range(4):
    add_mux(f'sat_mux{i}')

save_file(weights, 'model.safetensors')

def eval_xor(a, b):
    return int((a or b) and not (a and b))

def eval_fa(a, b, cin):
    x1 = eval_xor(a, b)
    s = eval_xor(x1, cin)
    c = int(a + b + cin >= 2)
    return s, c

def saturating_add(a, b):
    a_bits = [(a >> i) & 1 for i in range(4)]
    b_bits = [(b >> i) & 1 for i in range(4)]

    s_bits = []
    c = 0
    for i in range(4):
        s, c = eval_fa(a_bits[i], b_bits[i], c)
        s_bits.append(s)

    cout = c
    saturated = cout

    if saturated:
        result = 15
    else:
        result = sum(s_bits[i] << i for i in range(4))

    return result, saturated

print("Verifying 4-bit Saturating Adder...")
errors = 0
for a in range(16):
    for b in range(16):
        result, sat = saturating_add(a, b)
        true_sum = a + b
        if true_sum > 15:
            expected = 15
            exp_sat = 1
        else:
            expected = true_sum
            exp_sat = 0

        if result != expected or sat != exp_sat:
            errors += 1
            if errors <= 5:
                print(f"ERROR: {a}+{b} = {result} (sat={sat}), expected {expected} (sat={exp_sat})")

if errors == 0:
    print("All 256 test cases passed!")
else:
    print(f"FAILED: {errors} errors")

print("\nExamples:")
print("  7 + 5 = 12, saturated=0")
print("  10 + 10 = 15, saturated=1 (would be 20)")
print("  15 + 15 = 15, saturated=1 (would be 30)")

mag = sum(t.abs().sum().item() for t in weights.values())
print(f"\nMagnitude: {mag:.0f}")
print(f"Parameters: {sum(t.numel() for t in weights.values())}")