import numpy as np

def softmax(x):
  return 1 / (1 + np.exp(-x))

def deriv_softmax(x):
  fx = softmax(x)
  return fx * (1 - fx)

def mse_loss(y_true, y_pred):
  return ((y_true - y_pred) ** 2).mean()

class NeuralNetwork:
  def __init__(self):

    # Weights
    self.w1 = np.random.normal()
    self.w2 = np.random.normal()
    self.w3 = np.random.normal()
    self.w4 = np.random.normal()
    self.w5 = np.random.normal()
    self.w6 = np.random.normal()

    # Biases
    self.b1 = np.random.normal()
    self.b2 = np.random.normal()
    self.b3 = np.random.normal()

  def feedforward(self, x):

    h1 = softmax(self.w1 * x[0] + self.w2 * x[1] + self.b1)
    h2 = softmax(self.w3 * x[0] + self.w4 * x[1] + self.b2)
    o1 = softmax(self.w5 * h1 + self.w6 * h2 + self.b3)
    return o1

  def train(self, data, all_y_trues):
 
    lr = 0.1
    epochs = 1000

    for epoch in range(epochs):
      for x, y_true in zip(data, all_y_trues):
   
        sum_h1 = self.w1 * x[0] + self.w2 * x[1] + self.b1
        h1 = softmax(sum_h1)

        sum_h2 = self.w3 * x[0] + self.w4 * x[1] + self.b2
        h2 = softmax(sum_h2)

        sum_o1 = self.w5 * h1 + self.w6 * h2 + self.b3
        o1 = softmax(sum_o1)
        y_pred = o1

        d_L_d_ypred = 2 * (y_true - y_pred)

        # Neuron o1
        d_ypred_d_w5 = h1 * deriv_softmax(sum_o1)
        d_ypred_d_w6 = h2 * deriv_softmax(sum_o1)
        d_ypred_d_b3 = deriv_softmax(sum_o1)

        d_ypred_d_h1 = self.w5 * deriv_softmax(sum_o1)
        d_ypred_d_h2 = self.w6 * deriv_softmax(sum_o1)

        # Neuron h1
        d_h1_d_w1 = x[0] * deriv_softmax(sum_h1)
        d_h1_d_w2 = x[1] * deriv_softmax(sum_h1)
        d_h1_d_b1 = deriv_softmax(sum_h1)

        # Neuron h2
        d_h2_d_w3 = x[0] * deriv_softmax(sum_h2)
        d_h2_d_w4 = x[1] * deriv_softmax(sum_h2)
        d_h2_d_b2 = deriv_softmax(sum_h2)

        # Update weights and biases

        # Neuron h1
        self.w1 = self.w1 + lr * d_L_d_ypred * d_ypred_d_h1 * d_h1_d_w1
        self.w2 = self.w2 + lr * d_L_d_ypred * d_ypred_d_h1 * d_h1_d_w2
        self.b1 = self.b1 + lr * d_L_d_ypred * d_ypred_d_h1 * d_h1_d_b1

        # Neuron h2
        self.w3 =  self.w3 +  lr * d_L_d_ypred * d_ypred_d_h2 * d_h2_d_w3
        self.w4 =  self.w4 + lr * d_L_d_ypred * d_ypred_d_h2 * d_h2_d_w4
        self.b2 =  self.b2 + lr * d_L_d_ypred * d_ypred_d_h2 * d_h2_d_b2

        # Neuron o1
        self.w5 =  self.w5 + lr * d_L_d_ypred * d_ypred_d_w5
        self.w6 =  self.w6 + lr * d_L_d_ypred * d_ypred_d_w6
        self.b3 =  self.b3 + lr * d_L_d_ypred * d_ypred_d_b3

      if epoch % 100 == 0:
        y_preds = np.apply_along_axis(self.feedforward, 1, data)
        loss = mse_loss(all_y_trues, y_preds)
        print("Epoch %d loss: %.3f" % (epoch, loss))

data = np.array([[-2, -1],  [25, 6],   [17, 4],   [-15, -6]])
all_y_trues = np.array([1, 0, 0, 1])

network = NeuralNetwork()
network.train(data, all_y_trues)

test1 = np.array([-10, -7]) 
test2 = np.array([40, 20])  
print("test1: %.3f" % network.feedforward(test1)) 
print("test2: %.3f" % network.feedforward(test2)) 


