import torch
from safetensors.torch import save_file

weights = {}

# 2x2 Dadda Tree Multiplier
# Inputs: A1,A0, B1,B0 (4 inputs)
# Outputs: P3,P2,P1,P0 (4 outputs)
#
# Partial products:
#       A1B0  A0B0
#  A1B1 A0B1
# ----------------
#  P3   P2   P1   P0
#
# Dadda reduces column heights to specific sequence: 2,2,3,4,6,9,13...
# For 2x2, columns are already height 1 or 2, so minimal reduction needed.

# Input indices: A1=0, A0=1, B1=2, B0=3

def add_and(name, idx_a, idx_b, n_inputs):
    w = [0.0] * n_inputs
    w[idx_a] = 1.0
    w[idx_b] = 1.0
    weights[f'{name}.weight'] = torch.tensor([w], dtype=torch.float32)
    weights[f'{name}.bias'] = torch.tensor([-2.0], dtype=torch.float32)

def add_xor_direct(name, n_prev):
    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)

# Partial products (AND gates)
add_and('pp00', 1, 3, 4)  # A0 * B0 -> P0
add_and('pp01', 1, 2, 4)  # A0 * B1
add_and('pp10', 0, 3, 4)  # A1 * B0
add_and('pp11', 0, 2, 4)  # A1 * B1

# Column 1 (weight 2): pp01 + pp10 -> half adder
add_xor_direct('ha1_sum', 2)
weights['ha1_carry.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['ha1_carry.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# Column 2 (weight 4): pp11 + carry from column 1 -> half adder
add_xor_direct('ha2_sum', 2)
weights['ha2_carry.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['ha2_carry.bias'] = torch.tensor([-2.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

def dadda_mult(a1, a0, b1, b0):
    pp00 = a0 & b0
    pp01 = a0 & b1
    pp10 = a1 & b0
    pp11 = a1 & b1

    p0 = pp00

    # Half adder for column 1
    p1 = pp01 ^ pp10
    c1 = pp01 & pp10

    # Half adder for column 2
    p2 = pp11 ^ c1
    c2 = pp11 & c1

    p3 = c2

    return p3, p2, p1, p0

print("Verifying 2x2 Dadda tree multiplier...")
errors = 0
for a in range(4):
    for b in range(4):
        a1, a0 = (a >> 1) & 1, a & 1
        b1, b0 = (b >> 1) & 1, b & 1
        p3, p2, p1, p0 = dadda_mult(a1, a0, b1, b0)
        result = p3*8 + p2*4 + p1*2 + p0
        expected = a * b
        if result != expected:
            errors += 1
            print(f"ERROR: {a}*{b} = {result}, expected {expected}")

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

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