import torch
from safetensors.torch import save_file

weights = {}

# Canalizing Function
# Inputs: x2, x1, x0
# Output: y = NOT(x0) AND (x1 OR x2)
#
# Canalizing property: If x0 = 1, output is 0 regardless of x1, x2.
# x0 is the "canalizing input" with canalizing value 1 and canalized output 0.
#
# When x0 = 0, the function becomes y = x1 OR x2.
#
# This is NOT linearly separable (requires 2 layers).

# Layer 1: Compute intermediate values
# not_x0 = NOT(x0)
weights['not_x0.weight'] = torch.tensor([[0.0, 0.0, -1.0]], dtype=torch.float32)
weights['not_x0.bias'] = torch.tensor([0.0], dtype=torch.float32)

# or_x1_x2 = x1 OR x2 (inputs are [x2, x1, x0], so x1 is index 1, x2 is index 0)
weights['or_x1_x2.weight'] = torch.tensor([[1.0, 1.0, 0.0]], dtype=torch.float32)
weights['or_x1_x2.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# Layer 2: y = AND(not_x0, or_x1_x2)
weights['y.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['y.bias'] = torch.tensor([-2.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

def canalizing(x2, x1, x0):
    inp = torch.tensor([float(x2), float(x1), float(x0)])

    # Layer 1
    not_x0 = int((inp @ weights['not_x0.weight'].T + weights['not_x0.bias'] >= 0).item())
    or_x1_x2 = int((inp @ weights['or_x1_x2.weight'].T + weights['or_x1_x2.bias'] >= 0).item())

    # Layer 2
    l1 = torch.tensor([float(not_x0), float(or_x1_x2)])
    y = int((l1 @ weights['y.weight'].T + weights['y.bias'] >= 0).item())

    return y

def reference(x2, x1, x0):
    if x0 == 1:
        return 0
    return 1 if (x1 or x2) else 0

print("Verifying Canalizing Function...")
errors = 0
for i in range(8):
    x2, x1, x0 = (i >> 2) & 1, (i >> 1) & 1, i & 1
    result = canalizing(x2, x1, x0)
    expected = reference(x2, x1, x0)
    if result != expected:
        errors += 1
        print(f"ERROR: ({x2},{x1},{x0}) -> {result}, expected {expected}")

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

print("\nTruth Table:")
print("x2 x1 x0 | y | Canalizing?")
print("-" * 30)
for i in range(8):
    x2, x1, x0 = (i >> 2) & 1, (i >> 1) & 1, i & 1
    y = canalizing(x2, x1, x0)
    canal = "x0=1 -> y=0" if x0 == 1 else ""
    print(f" {x2}  {x1}  {x0}  | {y} | {canal}")

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())}")
print(f"Neurons: {len([k for k in weights.keys() if 'weight' in k])}")