#!/usr/bin/python3

import sys, os
import numpy as np

class CNN:

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#~~~~~~~~~~~~~~~~~~~~~~~~~~ INITIALIZING PARAMETERS ~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#

  # Load MNIST dataset.
  def load_MNIST(self):
    sys.stderr = open(os.devnull, 'w') # Disable stderr.
    from keras.datasets import mnist
    (self.X_train, self.Y_train), (self.X_test, self.Y_test) = mnist.load_data()
    self.num_trains = self.X_train.shape[0]
    self.num_tests = self.X_test.shape[0]
    self.X_train = self.X_train.astype('float32') / 255
    self.X_test = self.X_test.astype('float32') / 255
    self.Y_train = self.convert_to_one_hot(self.Y_train)
    self.Y_test = self.convert_to_one_hot(self.Y_test)
    sys.stderr = sys.__stderr__ # Enable stderr.

  # Convert labels to one-hot (10 x 1) vectors.
  def convert_to_one_hot(self, labels):
      labels_vecs = np.zeros((labels.size, 10, 1))
      labels_vecs[np.arange(labels.size), labels] = 1
      return labels_vecs

  # Initialize an array of a certain shape so that its values are uniformly
  # distributed on [-bound, bound].
  def initialize_uniformly(self, shape, bound):
    return self.scale_to_bound(np.random.rand(*shape), bound)

  # Given an array of values uniformly distributed on [0, 1], scale these values
  # so that they are uniformly distributed on [-bound, bound].
  def scale_to_bound(self, arr, bound):
    return arr * 2 * bound - bound

  def __init__(self):

    ###########################################
    #####           Load Inputs           #####
    ###########################################
    self.X_train = None
    self.Y_train = None
    self.num_trains = None
    self.X_test = None
    self.Y_test = None
    self.num_tests = None
    self.load_MNIST()

    # Input (image of digits) and true output (label) holders.
    self.X = np.empty(self.X_train[0].shape)
    self.Y = np.empty(self.Y_train[0].shape)
    self.X = np.expand_dims(self.X, axis=0)

    # Learning rates
    self.alpha = 1.0
    self.batch_size = 30
    self.batches = int(self.num_trains / self.batch_size)
    self.epochs = 10

    ###########################################
    #####      Convolutional Layer 1      #####
    ###########################################
    # (1 x 6) kernels of size (5 x 5).
    K1_shape = (1, 6, 5, 5)
    K1_bound = np.sqrt(6 / ((1 + 6) * (5 * 5)))
    self.K1 = self.initialize_uniformly(K1_shape, K1_bound)
    self.delta_K1 = np.zeros(self.K1.shape)
    self.batch_delta_K1 = np.zeros(self.K1.shape)

    # 6 biases.
    self.B1 = np.zeros((6, 1))
    self.delta_B1 = np.zeros(self.B1.shape)
    self.batch_delta_B1 = np.zeros(self.B1.shape)

    # 6 convolution outputs of size (24 x 24), before and after sigmoid.
    self.C1 = np.zeros((6, 24, 24))
    self.C1_sigma = np.zeros(self.C1.shape)
    self.delta_C1 = np.zeros(self.C1.shape)
    self.delta_C1_sigma = np.zeros(self.C1_sigma.shape)

    ###########################################
    #####         Pooling Layer 1         #####
    ###########################################
    # 6 sampling outputs of size (12 x 12).
    self.S1 = np.zeros((6, 12, 12))
    self.delta_S1 = np.zeros(self.S1.shape)

    ###########################################
    #####      Convolutional Layer 2      #####
    ###########################################
    # (6 x 12) kernels of size (5 x 5).
    K2_shape = (6, 12, 5, 5)
    K2_bound = np.sqrt(6 / ((6 + 12) * (5 * 5)))
    self.K2 = self.initialize_uniformly(K2_shape, K2_bound)
    self.delta_K2 = np.zeros(self.K2.shape)
    self.batch_delta_K2 = np.zeros(self.K2.shape)

    # 12 biases.
    self.B2 = np.zeros((12, 1))
    self.delta_B2 = np.zeros(self.B2.shape)
    self.batch_delta_B2 = np.zeros(self.B2.shape)

    # 12 convolution outputs of size (8 x 8), before and after sigmoid.
    self.C2 = np.zeros((12, 8, 8))
    self.C2_sigma = np.zeros(self.C2.shape)
    self.delta_C2 = np.zeros(self.C2.shape)
    self.delta_C2_sigma = np.zeros(self.C2_sigma.shape)

    ###########################################
    #####         Pooling Layer 2         #####
    ###########################################
    # 12 sampling outputs of size (4 x 4).
    self.S2 = np.zeros((12, 4, 4))
    self.delta_S2 = np.zeros(self.S2.shape)

    ###########################################
    #####      Fully Connected Layer      #####
    ###########################################
    # Vectorization of 12 matrices of size (4 x 4) to a vector of size 192.
    self.V = np.zeros((192, 1))
    self.delta_V = np.zeros(self.V.shape)

    # 10 weights of size 192.
    W_shape = (10, 192)
    W_bound = np.sqrt(6 / (192 + 10))
    self.W = self.initialize_uniformly(W_shape, W_bound)
    self.delta_W = np.zeros(self.W.shape)
    self.batch_delta_W = np.zeros(self.W.shape)

    # 10 biases.
    self.B = np.zeros((10, 1))
    self.delta_B = np.zeros(self.B.shape)
    self.batch_delta_B = np.zeros(self.B.shape)

    # 10 output classes, before and after sigmoid.
    self.Y_hat = np.zeros((10, 1))
    self.Y_hat_sigma = np.zeros(self.Y_hat.shape)
    self.delta_Y_hat = np.zeros(self.Y_hat.shape)
    self.delta_Y_hat_sigma = np.zeros(self.delta_Y_hat.shape)

    # Loss.
    self.L = 0.0
    self.batch_L = 0.0

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~ UNIVERSAL FUNCTIONS ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#

  ###########################################
  #####           Correlation           #####
  #####      (padding is not used)      #####
  ###########################################
  def correlation(self, I, K, O):
    # REFERENCE FORMULA: (5)
    # Verify the shape requirements for correlation with no padding.
    x_i, y_i = I.shape
    x_k, y_k = K.shape
    x_o, y_o = O.shape
    assert(x_o == x_i - x_k + 1)
    assert(y_o == y_i - y_k + 1)

    # REFERENCE FORMULA: (1)
    # Iterate through output matrix.
    for i in range(x_o):
      for j in range(y_o):
        O[i][j] = 0
        # Iterate through kernel.
        for u in range(x_k):
          for v in range(y_k):
            O[i][j] += I[i + u][j + v] * K[u][v]

  def correlation_wrapper_feedfrwd(self, all_I, all_K, all_B, all_O):
    num_I = all_I.shape[0]
    num_K = (all_K.shape[0], all_K.shape[1])
    num_B = all_B.shape[0]
    num_O = all_O.shape[0]

    assert(num_I == num_K[0])
    assert(num_O == num_K[1] == num_B)
    assert(all_B.shape[1] == 1)

    for p in range(num_I):
      for q in range(num_O):
        # Perform convolution.
        self.correlation(all_I[p], all_K[p][q], all_O[q])
        # Add bias.
        all_O[q] += all_B[q]

  def correlation_wrapper_backprop(self, all_I, all_K, all_O):
    num_I = all_I.shape[0]
    num_K = all_K.shape[0]
    num_O = (all_O.shape[0], all_O.shape[1])

    assert(num_I == num_O[0])
    assert(num_K == num_O[1])

    for p in range(num_I):
      for q in range(num_K):
        # Perform convolution.
        self.correlation(all_I[p], all_K[p], all_O[p][q])

  ###########################################
  #####           Convolution           #####
  #####        (padding is used)        #####
  ###########################################
  def convolution(self, I, K, O):

    # REFERENCE FORMULA: (6)
    # Verify the shape requirements for correlation with no padding.
    x_i, y_i = I.shape
    x_k, y_k = K.shape
    x_o, y_o = O.shape
    assert(x_o == x_i + x_k - 1)
    assert(y_o == y_i + y_k - 1)

    # REFERENCE FORMULA: (3) + (4)
    # Iterate through output matrix.
    for i in range(x_o):
      for j in range(y_o):
        O[i][j] = 0
        # Iterate through kernel.
        for u in range(x_k):
          for v in range(y_k):
            # Only care about indices within the bounds of input matrix
            # (because those outside the bounds are 0 anyway).
            if ((0 <= i - u) and (i - u <= x_i - 1) and
                (0 <= j - v) and (j - v <= y_i - 1)):
              O[i][j] += I[i - u][j - v] * K[u][v]

  def convolution_wrapper_backprop(self, all_I, all_K, all_O):

    # REFERENCE FORMULA: (78)
    num_I = all_I.shape[0]
    num_K = (all_K.shape[0], all_K.shape[1])
    num_O = all_O.shape[0]
    assert(num_I == num_K[1])
    assert(num_O == num_K[0])

    for p in range(num_O):
      all_O[p].fill(0)
      temp_O = np.empty(all_O[p].shape)
      for q in range(num_I):
        self.convolution(all_I[p], all_K[p][q], temp_O)
        all_O[p] += temp_O

  ###########################################
  #####             Sigmoid             #####
  ###########################################
  def sigmoid(self, I, O):

    # Verify the shapes are consistent.
    assert(O.shape == I.shape)

    # REFERENCE FORMULA: (7)
    np.copyto(O, 1 / (1 + np.exp(-I)))

  ###########################################
  #####             Pooling             #####
  ###########################################
  def pooling(self, I, O):

    # Ensure the side length of output is half the side length of the input.
    # Assume the side length of the input is even.
    x_i, y_i = I.shape
    x_o, y_o = O.shape
    assert(x_i % 2 == 0 and y_i % 2 == 0)
    assert(x_i == 2 * x_o)
    assert(y_i == 2 * y_o)

    # REFERENCE FORMULA: (12) + (16)
    # Iterate through the output image.
    for i in range(x_o):
      for j in range(y_o):
        O[i][j] = 0
        # Iterate through the pooling patch.
        for u in range(1):
          for v in range(1):
            O[i][j] += 0.25 * (I[2*i + u][2*j + v])

  def pooling_wrapper(self, all_I, all_O):

    num_I = all_I.shape[0]
    num_O = all_O.shape[0]
    assert(num_I == num_O)

    for p in range(num_I):
      self.pooling(all_I[p], all_O[p])

  def unpooling(self, I, O):

    # Ensure the side length of output is twice the side length of the input.
    x_i, y_i = I.shape
    x_o, y_o = O.shape
    assert(x_o == 2 * x_i)
    assert(y_o == 2 * y_i)

    # REFERENCE FORMULA: (55) + (83)
    # Iterate through the output image.
    for i in range(x_o):
      for j in range(y_o):
        O[i][j] = 0.25 * I[int(i / 2)][int(j / 2)]

  def unpooling_wrapper(self, all_I, all_O):

    num_I = all_I.shape[0]
    num_O = all_O.shape[0]
    assert(num_I == num_O)

    for p in range(num_I):
      self.unpooling(all_I[p], all_O[p])

  ###########################################
  #####          Vectorization          #####
  ###########################################
  def vectorize(self, I, O):

    # Verify that the output is a vector.
    assert(len(O.shape) == 2)
    assert(O.shape[1] == 1)

    # Verify that the number of cells in the input matrix is equal to the
    # length of the output vector.
    assert(np.prod(I.shape) == O.shape[0])

    np.copyto(O, I.reshape(O.shape))

  def unvectorize(self, I, O):

    # Verify that the input is a vector.
    assert(len(I.shape) == 2)
    assert(I.shape[1] == 1)

    # Verify that the number of cells in the output matrix is equal to the
    # length of the input vector.
    assert(np.prod(O.shape) == I.shape[0])

    np.copyto(O, I.reshape(O.shape))

  ###########################################
  #####         Fully Connected         #####
  ###########################################
  def fully_connected(self, I, W, B, O):

    # Verify the shapes are consistent.
    assert(O.shape[0] == W.shape[0] == B.shape[0])
    assert(O.shape[1] == I.shape[1] == B.shape[1])
    assert(W.shape[1] == I.shape[0])

    np.copyto(O, np.dot(W, I) + B)

  ###########################################
  #####          Loss Function          #####
  ###########################################
  def loss(self, truth, prediction):

    # Verify the shapes are consistent.
    assert(truth.shape == prediction.shape)

    self.L = 0.5 * np.sum(np.square((truth - prediction)))
    self.batch_L += self.L / self.batch_size

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ FEEDFORWARD ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  def feed_forward(self, X, Y):
    assert(self.X[0].shape == X.shape)
    self.X[0] = X
    assert(self.Y_hat.shape == Y.shape)
    self.Y = Y

    ###########################################
    ######      Convolution Layer 1      ######
    ###########################################
    self.correlation_wrapper_feedfrwd(self.X, self.K1, self.B1, self.C1)
    self.sigmoid(self.C1, self.C1_sigma)

    ###########################################
    ######        Pooling Layer 1        ######
    ###########################################
    self.pooling_wrapper(self.C1_sigma, self.S1)

    ###########################################
    ######      Convolution Layer 2      ######
    ###########################################
    self.correlation_wrapper_feedfrwd(self.S1, self.K2, self.B2, self.C2)
    self.sigmoid(self.C2, self.C2_sigma)

    ###########################################
    ######        Pooling Layer 2        ######
    ###########################################
    self.pooling_wrapper(self.C2_sigma, self.S2)

    ###########################################
    ######         Vectorization         ######
    ###########################################
    self.vectorize(self.S2, self.V)

    ###########################################
    ######     Fully Connected Layer     ######
    ###########################################
    self.fully_connected(self.V, self.W, self.B, self.Y_hat)
    self.sigmoid(self.Y_hat, self.Y_hat_sigma)

    ###########################################
    #####          Loss Function          #####
    ###########################################
    self.loss(self.Y, self.Y_hat_sigma)

  def back_propagation(self):
    ###########################################
    #####          Loss Function          #####
    ###########################################
    self.delta_Y_hat_sigma = self.Y_hat_sigma - self.Y
    self.delta_Y_hat = np.multiply(
      self.delta_Y_hat_sigma,
      np.multiply(self.Y_hat_sigma, 1 - self.Y_hat_sigma)
    )

    ###########################################
    ######     Fully Connected Layer     ######
    ###########################################
    self.delta_W = self.delta_Y_hat.dot(self.V.T)
    self.delta_B = self.delta_Y_hat

    ###########################################
    ######         Vectorization         ######
    ###########################################
    self.delta_V = self.W.T.dot(self.delta_Y_hat)

    ###########################################
    ######        Pooling Layer 2        ######
    ###########################################
    self.unvectorize(self.delta_V, self.delta_S2)

    ###########################################
    ######      Convolution Layer 2      ######
    ###########################################
    self.unpooling_wrapper(self.delta_S2, self.delta_C2_sigma)
    self.delta_C2 = np.multiply(
      self.delta_C2_sigma,
      np.multiply(self.C2_sigma, 1 - self.C2_sigma)
    )
    self.correlation_wrapper_backprop(self.S1, self.delta_C2, self.delta_K2)
    for q in range(self.delta_B2.shape[0]):
      self.delta_B2[q] = np.sum(self.delta_C2[q])

    ###########################################
    ######        Pooling Layer 1        ######
    ###########################################
    self.convolution_wrapper_backprop(self.delta_C2, self.K2, self.delta_S1)

    ###########################################
    ######      Convolution Layer 2      ######
    ###########################################
    self.unpooling_wrapper(self.delta_S1, self.delta_C1_sigma)
    self.delta_C1 = np.multiply(
      self.delta_C1_sigma,
      np.multiply(self.C1_sigma, 1 - self.C1_sigma)
    )
    self.correlation_wrapper_backprop(self.X, self.delta_C1, self.delta_K1)
    for p in range(self.delta_B1.shape[0]):
      self.delta_B1[p] = np.sum(self.delta_C1[p])

    self.batch_delta_K1 += self.delta_K1 / self.batch_size
    self.batch_delta_B1 += self.delta_B1 / self.batch_size
    self.batch_delta_K2 += self.delta_K2 / self.batch_size
    self.batch_delta_B2 += self.delta_B2 / self.batch_size
    self.batch_delta_W  += self.delta_W  / self.batch_size
    self.batch_delta_B  += self.delta_B  / self.batch_size

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~ UPDATING PARAMETERS ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  def update_parameters(self):
    self.K1 -= self.alpha * self.batch_delta_K1
    self.B1 -= self.alpha * self.batch_delta_B1
    self.K2 -= self.alpha * self.batch_delta_K2
    self.B2 -= self.alpha * self.batch_delta_B2
    self.W  -= self.alpha * self.batch_delta_W
    self.B  -= self.alpha * self.batch_delta_B

    self.batch_delta_K1 = 0.0
    self.batch_delta_B1 = 0.0
    self.batch_delta_K2 = 0.0
    self.batch_delta_B2 = 0.0
    self.batch_delta_W  = 0.0
    self.batch_delta_B  = 0.0
    self.batch_L = 0.0

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~ EVALUATE ON TEST DATA ~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  def evaluate_test(self):
    accuracy = 0.0
    test_size = 30
    for i in range(test_size):
      cur = np.random.randint(self.num_tests)
      myCNN.feed_forward(
        myCNN.X_test[cur],
        myCNN.Y_test[cur]
      )
      prediction = np.argmax(myCNN.Y_hat_sigma)
      truth = np.argmax(myCNN.Y)
      accuracy += (prediction == truth)
      print(">> Tested {} / {}".format(i + 1, test_size),
            end=('\r' if i != test_size - 1 else ''))
    print(". Test accuracy:", accuracy / test_size)

myCNN = CNN()
print("Number of epochs", myCNN.epochs)
print("Number of batches", myCNN.batches)
print("Batch size", myCNN.batch_size)

for e in range(myCNN.epochs):
  print("EPOCH {} / {} ========================================".format(
        e + 1, myCNN.epochs))
  for b in range(myCNN.batches):
    print("BATCH {} / {} ----------".format(b + 1, myCNN.batches))
    for i in range(myCNN.batch_size):
      cur = b * myCNN.batch_size + i
      myCNN.feed_forward(
        myCNN.X_train[cur],
        myCNN.Y_train[cur]
      )
      myCNN.back_propagation()
      print(">> Trained {} / {}. Loss: {}".format(
            i + 1, myCNN.batch_size, myCNN.L), end='\r')
      sys.stdout.flush()
    print()
    myCNN.update_parameters()
  myCNN.evaluate_test()
    
