import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import itertools as it
 
%matplotlib inline

def sigmoid(x):
    """Sigmoid function"""
    return 1.0 / (1.0 + np.exp(-x))

# Plot The sigmoid function
xs = np.linspace(-10, 10, num=100, dtype=np.float32)
activation = sigmoid(xs)
 
fig = plt.figure(figsize=(6,4))
plt.plot(xs, activation)
plt.plot(0,.5,'ro')
 
plt.grid(True, which='both')
plt.axhline(y=0, color='y')
plt.axvline(x=0, color='y')
plt.ylim([-0.1, 1.15])

(-0.1, 1.15)

# Logic gate function
def logic_gate(w1, w2, b, x1, x2):
  ''' logic_gate is a function which returns the results of
        taking two args and  (hopefully) acts like a logic gate (and/or/not/etc.).
        its behavior is determined by w1,w2,b. '''
  return sigmoid(w1 * x1 + w2 * x2 + b)

# Test function that takes a function with two arguments
def test_gate(w1, w2, b):
  for x1 in range(2):
    for x2 in range(2):
      print("{}, {}: {}".format(x1, x2, np.round(logic_gate(w1, w2, b, x1, x2))))

or_gate = test_gate(20, 20, -10)

0, 0: 0.0
0, 1: 1.0
1, 0: 1.0
1, 1: 1.0

# Fill in the w1, w2, and b parameters such that the truth table matches
# and_gate = test_gate(...)

# Fill in the w1, w2, and b parameters such that the truth table matches
# nor_gate = test_gate(...)

# Fill in the w1, w2, and b parameters such that the truth table matches
# nand_gate = test_gate(...)

# Define LogicGate Model using tf.Module
class LogicGate(tf.Module):
    def __init__(self):
        super().__init__()
        self.built = False  # Track if model is initialized
 
    def __call__(self, x, train=True):
        # Initialize weights and bias on first call
        if not self.built:
            input_dim = x.shape[-1]  # Number of input features
            self.w = tf.Variable(tf.random.normal([input_dim, 1]), name="weights")
            self.b = tf.Variable(tf.zeros([1]), name="bias")
            self.built = True
 
        # Compute logits: z = Wx + b
        z = tf.add(tf.matmul(x, self.w), self.b)
        return tf.sigmoid(z)  # Apply sigmoid
 
# Loss function (Mean Squared Error)
def compute_loss(y_pred, y_true):
    return tf.reduce_mean(tf.square(y_pred - y_true))
 
# Training function
def train_model(model, x_train, y_train, learning_rate=0.5, epochs=5000):
  # Iterate over the training data
  for epoch in range(epochs):
    with tf.GradientTape() as tape:
      y_pred = model(x_train)  # Forward pass
      loss = compute_loss(y_pred, y_train)
 
    # Update the parameters with respect to the gradient calculations
    grads = tape.gradient(loss, model.variables)
    for g,v in zip(grads, model.variables):
      v.assign_sub(learning_rate * g)
 
    # Print progress every 1000 epochs
    if epoch % 1000 == 0:
      acc = compute_accuracy(model, x_train, y_train)
      print(f"Epoch {epoch}, Loss: {loss.numpy():.4f}, Accuracy: {acc:.4f}")
 
# Accuracy function
def compute_accuracy(model, x, y_true):
    y_pred = model(x, train=False)
    y_pred_rounded = tf.round(y_pred)
    correct = tf.equal(y_pred_rounded, y_true)
    return tf.reduce_mean(tf.cast(correct, tf.float32)).numpy()

 
# Prepare AND gate dataset
and_table = np.array([[0, 0, 0],
                      [1, 0, 0],
                      [0, 1, 0],
                      [1, 1, 1]], dtype=np.float32)
 
x_train = and_table[:, :2]  # Inputs: x1, x2
y_train = and_table[:, 2:]  # Labels: y
 
# Initialize and train model
model = LogicGate()
train_model(model, x_train, y_train)
 
# Evaluate and print results
w1, w2 = model.w.numpy().flatten()
b = model.b.numpy().flatten()[0]
print(f"\nLearned weight for w1: {w1}")
print(f"Learned weight for w2: {w2}")
print(f"Learned bias: {b}\n")
 
# Test model predictions
y_pred = model(x_train, train=False).numpy().round().astype(np.uint8)
print("Predicted Truth Table:")
print(np.column_stack((and_table[:, :2], y_pred)))

Epoch 0, Loss: 0.3312, Accuracy: 0.5000
Epoch 1000, Loss: 0.0126, Accuracy: 1.0000
Epoch 2000, Loss: 0.0058, Accuracy: 1.0000
Epoch 3000, Loss: 0.0037, Accuracy: 1.0000
Epoch 4000, Loss: 0.0027, Accuracy: 1.0000
 
Learned weight for w1: 5.733343124389648
Learned weight for w2: 5.733343124389648
Learned bias: -8.691651344299316
 
Predicted Truth Table:
[[0. 0. 0.]
 [1. 0. 0.]
 [0. 1. 0.]
 [1. 1. 1.]]

# Make sure you have or_gate, nand_gate, and and_gate working from above
#def xor_gate(a, b):
#    c = or_gate(a, b)
#    d = nand_gate(a, b)
#    return and_gate(c, d)
#test(xor_gate)