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

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

xs = np.linspace(-10, 10, num=100, dtype=np.float32)
plt.plot(xs, sigmoid(xs))
plt.plot(0, .5, 'ro')

def logic_gate(w1, w2, b, x1, x2):
  return sigmoid(w1 * x1 + w2 * x2 + b)

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 w1, w2, b so the truth table matches:
# and_gate = test_gate(...)
# nor_gate = test_gate(...)
# nand_gate = test_gate(...)

class LogicGate(tf.Module):
    def __init__(self):
        super().__init__()
        self.built = False

    def __call__(self, x, train=True):
        if not self.built:
            input_dim = x.shape[-1]
            self.w = tf.Variable(tf.random.normal([input_dim, 1]))
            self.b = tf.Variable(tf.zeros([1]))
            self.built = True
        z = tf.add(tf.matmul(x, self.w), self.b)
        return tf.sigmoid(z)

def compute_loss(y_pred, y_true):
    return tf.reduce_mean(tf.square(y_pred - y_true))

def train_model(model, x_train, y_train, learning_rate=0.5, epochs=5000):
  for epoch in range(epochs):
    with tf.GradientTape() as tape:
      y_pred = model(x_train)
      loss = compute_loss(y_pred, y_train)
    grads = tape.gradient(loss, model.variables)
    for g, v in zip(grads, model.variables):
      v.assign_sub(learning_rate * g)

def compute_accuracy(model, x, y_true):
    y_pred = tf.round(model(x, train=False))
    correct = tf.equal(y_pred, y_true)
    return tf.reduce_mean(tf.cast(correct, tf.float32)).numpy()

and_table = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 1]], dtype=np.float32)
x_train = and_table[:, :2]
y_train = and_table[:, 2:]

model = LogicGate()
train_model(model, x_train, y_train)

Epoch 0, Loss: 0.3312, Accuracy: 0.5000
Epoch 1000, Loss: 0.0126, Accuracy: 1.0000
Epoch 4000, Loss: 0.0027, Accuracy: 1.0000

Learned weight for w1: 5.733
Learned weight for w2: 5.733
Learned bias: -8.692

# def xor_gate(a, b):
#     c = or_gate(a, b)
#     d = nand_gate(a, b)
#     return and_gate(c, d)