from keras.datasets import mnist
import numpy as np
import pandas as pd


# Base class
class Layer:
    def __init__(self):
        self.input = None
        self.output = None

    # computes the output Y of a layer for a given input X
    def forward_propagation(self, input):
        raise NotImplementedError

    # computes dE/dX for a given dE/dY (and update parameters if any)
    def backward_propagation(self, output_error, learning_rate):
        raise NotImplementedError

class FCLayer(Layer):
    def __init__(self, input_size, output_size):
      self.weights = np.random.rand(input_size, output_size) - 0.5
      self.bias = np.random.rand(1, output_size) - 0.5
    
    def forward_propagation(self,input_data):
      self.input = input_data
      self.output = np.dot(self.input,self.weights)+self.bias
      return(self.output)

    ## compute dE/dw, de/dB for a given outputt erroor. Returns input error to be given t the next layer
    def backward_propagation(self,output_error, learning_rate):
      input_error = np.dot(output_error, self.weights.T)
      weights_error = np.dot(self.input.T, output_error)

      self.weights -= learning_rate * weights_error
      self.bias -= learning_rate * output_error
      return(input_error)


class ActivationLayer(Layer):
    def __init__(self):
      self.activation = lambda x: np.tanh(x)
      self.activation_prime = lambda x: 1-np.tanh(x)**2;
    
    def forward_propagation(self, input_data):
      self.input = input_data
      self.output = self.activation(self.input)
      return(self.output)
    
    def backward_propagation(self, output_error, learning_rate):
      return(self.activation_prime(self.input) * output_error)

# loss function and its derivative
def mse(y_true, y_pred):
    return np.mean(np.power(y_true - y_pred, 2));

def mse_prime(y_true, y_pred):
    return 2*(y_pred-y_true)/y_true.size;

class Network:
  def __init__(self):
    self.layers=[]
    self.loss=None
    self.loss_prime=None

  def add(self,layer):
    self.layers.append(layer)

  def use(self,loss,loss_prime):
    self.loss = loss
    self.loss_prime = loss_prime

  def predict(self,input_data):
    result=[]
    for i in range(len(input_data)):
      output=input_data[i]
      for layer in self.layers:
        output= layer.forward_propagation(output)
      result.append(output)
    return(result)

  def fit(self, x_train, y_train, epochs, learning_rate):
    samples = len(x_train)
    errors=[]

    for i in range(epochs):
      err = 0
      for j in range(samples):
        output = x_train[j]
        for layer in self.layers:
          output = layer.forward_propagation(output)
        
        err += self.loss(y_train[j], output)

        error = self.loss_prime(y_train[j], output)
        for layer in reversed(self.layers):
          error = layer.backward_propagation(error, learning_rate)
      errors.append(err/samples)

      print('epoch %d/%d   error=%f' % (i+1, epochs, err))
    return(errors)



x_train = np.array([[[0,0]], [[0,1]], [[1,0]], [[1,1]]])
y_train = np.array([[[0]], [[1]], [[1]], [[0]]])


net = Network()
net.add(FCLayer(2,5))
net.add(ActivationLayer())
net.add(FCLayer(5,1))
net.add(ActivationLayer())
net.use(mse, mse_prime)

err = net.fit(x_train, y_train, epochs=1000, learning_rate=0.1)

pd.DataFrame(err).plot()

out = net.predict(x_train)
print(x_train, out)

###### Digits recognitin
from keras.datasets import mnist
from keras.utils import np_utils

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(x_train.shape[0], 1, 28*28).astype('float32')/255
y_train = np_utils.to_categorical(y_train)

x_test = x_test.reshape(x_test.shape[0], 1, 28*28).astype('float32')/255
y_test = np_utils.to_categorical(y_test)

y_train[0]

net = Network()
net.add(FCLayer(28*28, 100))
net.add(ActivationLayer())
net.add(FCLayer(100, 50)) 
net.add(ActivationLayer())
net.add(FCLayer(50, 10)) 
net.add(ActivationLayer())
net.use(mse, mse_prime)

errors = net.fit(x_train[0:5000], y_train[0:5000], epochs=35, learning_rate=0.1)

errors=[]
for i in range(1000):
  out=sum((net.predict(x_test[i]) - y_test[i])[0][0])
  errors.append(0 if out<0.5 else 1)


np.mean(errors)