import torch
from safetensors.torch import save_file

weights = {}

# 2x2 Array Multiplier
# Uses regular array structure of AND gates and adders
# Inputs: A1,A0, B1,B0 (4 inputs)
# Outputs: P3,P2,P1,P0 (4 outputs)

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):
    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_ha_carry(name):
    weights[f'{name}.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
    weights[f'{name}.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# Row 0: A[1:0] * B0
add_and('r0c0', 1, 3, 4)  # A0*B0 -> P0
add_and('r0c1', 0, 3, 4)  # A1*B0

# Row 1: A[1:0] * B1, shifted and added
add_and('r1c0', 1, 2, 4)  # A0*B1
add_and('r1c1', 0, 2, 4)  # A1*B1

# Array cell (0,1): r0c1 + r1c0 -> half adder
add_xor_direct('cell01_sum')
add_ha_carry('cell01_carry')

# Array cell (1,1): r1c1 + carry -> half adder
add_xor_direct('cell11_sum')
add_ha_carry('cell11_carry')

save_file(weights, 'model.safetensors')

def array_mult(a1, a0, b1, b0):
    # Partial products
    r0c0 = a0 & b0  # P0
    r0c1 = a1 & b0
    r1c0 = a0 & b1
    r1c1 = a1 & b1

    # Array addition
    p0 = r0c0
    p1 = r0c1 ^ r1c0
    c1 = r0c1 & r1c0
    p2 = r1c1 ^ c1
    p3 = r1c1 & c1

    return p3, p2, p1, p0

print("Verifying 2x2 array 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 = array_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())}")