import torch
from safetensors.torch import save_file

# Absolute value of 4-bit 2's complement number
# Input: a3 (sign), a2, a1, a0
# Output: o3, o2, o1, o0
#
# For positive (a3=0): output = input
# For negative (a3=1): output = ~input + 1 (2's complement)
#
# 2's complement bit formulas:
#   o0 = a0 (always)
#   o1 = ~a1 XOR ~a0 (when negative)
#   o2 = ~a2 XOR (~a1 AND ~a0) (when negative)
#   o3 = ~a3 XOR (~a2 AND ~a1 AND ~a0) = carry (when negative, and ~a3=0 for a3=1)

weights = {}

# === DIRECT OUTPUTS ===
# o0 = a0 (always)
weights['o0.weight'] = torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32)
weights['o0.bias'] = torch.tensor([-0.5], dtype=torch.float32)

# o3 = 1 only when input is 1000 (-8), i.e., when carry propagates all the way
# o3 = a3 AND NOT(a2) AND NOT(a1) AND NOT(a0)
weights['o3.weight'] = torch.tensor([[1.0, -1.0, -1.0, -1.0]], dtype=torch.float32)
weights['o3.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# === LAYER 1 (from raw inputs a3, a2, a1, a0) ===

# NOT gates
weights['l1.not_a0.weight'] = torch.tensor([[0.0, 0.0, 0.0, -1.0]], dtype=torch.float32)
weights['l1.not_a0.bias'] = torch.tensor([0.0], dtype=torch.float32)

weights['l1.not_a1.weight'] = torch.tensor([[0.0, 0.0, -1.0, 0.0]], dtype=torch.float32)
weights['l1.not_a1.bias'] = torch.tensor([0.0], dtype=torch.float32)

weights['l1.not_a2.weight'] = torch.tensor([[0.0, -1.0, 0.0, 0.0]], dtype=torch.float32)
weights['l1.not_a2.bias'] = torch.tensor([0.0], dtype=torch.float32)

# XOR(~a1, ~a0) = o1 for negative path
# Components: OR(~a1, ~a0), NAND(~a1, ~a0)
# OR(~a1, ~a0) = NOR(a1, a0) inverted... actually OR(~a1,~a0) = NAND(a1,a0)
weights['l1.nand10.weight'] = torch.tensor([[0.0, 0.0, -1.0, -1.0]], dtype=torch.float32)
weights['l1.nand10.bias'] = torch.tensor([1.0], dtype=torch.float32)

# NAND(~a1, ~a0) = NOT(~a1 AND ~a0) = NOT(NOR(a1,a0)) = OR(a1,a0)
weights['l1.or10.weight'] = torch.tensor([[0.0, 0.0, 1.0, 1.0]], dtype=torch.float32)
weights['l1.or10.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# NOR(a1, a0) = ~a1 AND ~a0 = carry for bit 2
weights['l1.nor10.weight'] = torch.tensor([[0.0, 0.0, -1.0, -1.0]], dtype=torch.float32)
weights['l1.nor10.bias'] = torch.tensor([0.0], dtype=torch.float32)

# Positive path: pass through a1 when a3=0
weights['l1.o1_pos.weight'] = torch.tensor([[-1.0, 0.0, 1.0, 0.0]], dtype=torch.float32)
weights['l1.o1_pos.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# Positive path: pass through a2 when a3=0
weights['l1.o2_pos.weight'] = torch.tensor([[-1.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
weights['l1.o2_pos.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# a3 passthrough
weights['l1.a3.weight'] = torch.tensor([[1.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['l1.a3.bias'] = torch.tensor([-0.5], dtype=torch.float32)

# === LAYER 2 ===
# Inputs: [not_a0, not_a1, not_a2, nand10, or10, nor10, o1_pos, o2_pos, a3]

# XOR(~a1, ~a0) = nand10 AND or10  (since nand10 = OR(~a1,~a0), or10 = NAND(~a1,~a0))
# Wait, let me reconsider:
# XOR(~a1, ~a0) = (~a1 OR ~a0) AND NOT(~a1 AND ~a0)
#              = NAND(a1,a0) AND OR(a1,a0)
# nand10 = NAND(a1,a0) ✓
# or10 = OR(a1,a0) ✓
weights['l2.xor_neg_10.weight'] = torch.tensor([[0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['l2.xor_neg_10.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# For XOR(~a2, nor10): need OR(~a2, nor10) AND NAND(~a2, nor10)
# We have not_a2 and nor10 from layer 1
weights['l2.or_nota2_nor10.weight'] = torch.tensor([[0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['l2.or_nota2_nor10.bias'] = torch.tensor([-1.0], dtype=torch.float32)

weights['l2.nand_nota2_nor10.weight'] = torch.tensor([[0.0, 0.0, -1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['l2.nand_nota2_nor10.bias'] = torch.tensor([1.0], dtype=torch.float32)

# Pass through o1_pos, o2_pos, a3
weights['l2.o1_pos.weight'] = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
weights['l2.o1_pos.bias'] = torch.tensor([-0.5], dtype=torch.float32)

weights['l2.o2_pos.weight'] = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0]], dtype=torch.float32)
weights['l2.o2_pos.bias'] = torch.tensor([-0.5], dtype=torch.float32)

weights['l2.a3.weight'] = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]], dtype=torch.float32)
weights['l2.a3.bias'] = torch.tensor([-0.5], dtype=torch.float32)

# === LAYER 3 ===
# Inputs: [xor_neg_10, or_nota2_nor10, nand_nota2_nor10, o1_pos, o2_pos, a3]

# XOR(~a2, nor10) = or_nota2_nor10 AND nand_nota2_nor10
weights['l3.xor_neg_2.weight'] = torch.tensor([[0.0, 1.0, 1.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['l3.xor_neg_2.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# o1_neg = xor_neg_10 AND a3
weights['l3.o1_neg.weight'] = torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 1.0]], dtype=torch.float32)
weights['l3.o1_neg.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# Pass through o1_pos, o2_pos, a3
weights['l3.o1_pos.weight'] = torch.tensor([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
weights['l3.o1_pos.bias'] = torch.tensor([-0.5], dtype=torch.float32)

weights['l3.o2_pos.weight'] = torch.tensor([[0.0, 0.0, 0.0, 0.0, 1.0, 0.0]], dtype=torch.float32)
weights['l3.o2_pos.bias'] = torch.tensor([-0.5], dtype=torch.float32)

weights['l3.a3.weight'] = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 1.0]], dtype=torch.float32)
weights['l3.a3.bias'] = torch.tensor([-0.5], dtype=torch.float32)

# === LAYER 4 ===
# Inputs: [xor_neg_2, o1_neg, o1_pos, o2_pos, a3]

# o2_neg = xor_neg_2 AND a3
weights['l4.o2_neg.weight'] = torch.tensor([[1.0, 0.0, 0.0, 0.0, 1.0]], dtype=torch.float32)
weights['l4.o2_neg.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# o1 = o1_pos OR o1_neg
weights['l4.o1.weight'] = torch.tensor([[0.0, 1.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
weights['l4.o1.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# Pass through o2_pos
weights['l4.o2_pos.weight'] = torch.tensor([[0.0, 0.0, 0.0, 1.0, 0.0]], dtype=torch.float32)
weights['l4.o2_pos.bias'] = torch.tensor([-0.5], dtype=torch.float32)

# === LAYER 5 ===
# Inputs: [o2_neg, o2_pos]
# o2 = o2_pos OR o2_neg
weights['l5.o2.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['l5.o2.bias'] = torch.tensor([-1.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

# Verification
def abs4(a3, a2, a1, a0):
    inp = torch.tensor([float(a3), float(a2), float(a1), float(a0)])

    # Direct outputs
    o0 = int((inp @ weights['o0.weight'].T + weights['o0.bias'] >= 0).item())
    o3 = int((inp @ weights['o3.weight'].T + weights['o3.bias'] >= 0).item())

    # Layer 1
    l1_keys = ['not_a0', 'not_a1', 'not_a2', 'nand10', 'or10', 'nor10', 'o1_pos', 'o2_pos', 'a3']
    l1 = {k: int((inp @ weights[f'l1.{k}.weight'].T + weights[f'l1.{k}.bias'] >= 0).item()) for k in l1_keys}
    l1_out = torch.tensor([float(l1[k]) for k in l1_keys])

    # Layer 2
    l2_keys = ['xor_neg_10', 'or_nota2_nor10', 'nand_nota2_nor10', 'o1_pos', 'o2_pos', 'a3']
    l2 = {k: int((l1_out @ weights[f'l2.{k}.weight'].T + weights[f'l2.{k}.bias'] >= 0).item()) for k in l2_keys}
    l2_out = torch.tensor([float(l2[k]) for k in l2_keys])

    # Layer 3
    l3_keys = ['xor_neg_2', 'o1_neg', 'o1_pos', 'o2_pos', 'a3']
    l3 = {k: int((l2_out @ weights[f'l3.{k}.weight'].T + weights[f'l3.{k}.bias'] >= 0).item()) for k in l3_keys}
    l3_out = torch.tensor([float(l3[k]) for k in l3_keys])

    # Layer 4
    l4_keys = ['o2_neg', 'o1', 'o2_pos']
    l4 = {k: int((l3_out @ weights[f'l4.{k}.weight'].T + weights[f'l4.{k}.bias'] >= 0).item()) for k in l4_keys}

    o1 = l4['o1']

    # Layer 5
    l5_inp = torch.tensor([float(l4['o2_neg']), float(l4['o2_pos'])])
    o2 = int((l5_inp @ weights['l5.o2.weight'].T + weights['l5.o2.bias'] >= 0).item())

    return o3, o2, o1, o0

print("Verifying absolutevalue4...")
errors = 0
for a in range(16):
    a3, a2, a1, a0 = (a >> 3) & 1, (a >> 2) & 1, (a >> 1) & 1, a & 1
    o3, o2, o1, o0 = abs4(a3, a2, a1, a0)
    result = 8*o3 + 4*o2 + 2*o1 + o0

    signed_val = a if a < 8 else a - 16
    expected = abs(signed_val)

    if result != expected:
        errors += 1
        print(f"ERROR: input={a:04b} ({signed_val:+d}), got {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}")