'''
Author: SJ2050
Date: 2022-01-16 17:16:10
LastEditTime: 2022-01-22 10:09:33
Version: v0.0.1
Description: Forward-Propagation and Back-Propagation algorithms of MLP using function.
Copyright © 2022 SJ2050
'''
import numpy as np
from sklearn.metrics import accuracy_score

# NN = {'nodes_size': [  ],
#       'layers': [   ],
#       'w': [   ],
#       'b': [   ]}

def sigmoid(x):
  y  = 1 / (1+np.exp(-x))
  return y

def sigmoid_derivative(y):
  return y * (1-y)

def initialize(nodes_size):
  NN                     = {} 
  NN['nodes_size']       = nodes_size
  total_layer_num        = len(nodes_size)
  NN['layers']           = [np.array([]) for _ in range(total_layer_num)]
  NN['w']                = []
  NN['b']                = []

  for i in range(0, total_layer_num-1):
    NN['w'].append(np.random.random((nodes_size[i], nodes_size[i+1]))*2-1)
    NN['b'].append(np.zeros(nodes_size[i+1]))

  return NN

def forward_propagation(NN, input_layer):
  layer_num = len(NN['nodes_size'])
  NN['layers'][0] = input_layer

  x = input_layer
  for i in range(0, layer_num-1):
    w   = NN['w'][i]
    b   = NN['b'][i]
    x   = sigmoid(np.matmul(x, w)+b)
    NN['layers'][i+1] = x

def backward_propagation(NN, input_layer, tags, eps = 0.01):
  layer_num = len(NN['nodes_size'])
  t = tags

  forward_propagation(NN, input_layer)
  D = []
  # compute delta
  for j in range(layer_num-1, 0, -1):
    i = j - 1
    x = NN['layers'][i]
    y = NN['layers'][j]
    if j == layer_num - 1:
      d = sigmoid_derivative(y)*(t-y)
    else:
      w = NN['w'][j]
      d = sigmoid_derivative(y)*np.matmul(d, w.T)

    D.insert(0, d)

  # update weights and biases
  for j in range(layer_num-1, 0, -1):
    i = j - 1
    x = NN['layers'][i]
    d = D[i]
    NN['w'][i] += eps*np.matmul(x.T, d) / x.shape[0]
    NN['b'][i] += eps*np.sum(d, axis=0) / x.shape[0]

def train(NN, input_layer, tags, eps, iter_num, eval_num, batch_size):
  print('Training...')
  for i in range(iter_num):
    # print(f'{i+1}th training of {iter_num}.')
    if (batch_size < 0):
      batch_input = input_layer
      batch_tags  = tags
    else:
      index       = np.random.randint(0, input_layer.shape[0], batch_size)
      batch_input = input_layer[index]
      batch_tags  = tags[index]
    backward_propagation(NN, batch_input, batch_tags, eps)

    if (i+1) % eval_num == 0:
      forward_propagation(NN, input_layer)
      loss, acc = evaluate(NN, tags)
      print(f'{i+1}th training of {iter_num}: Loss={loss}, acc={acc}.')
  print('Training finished!')
  return NN['layers'][-1]

def predict(NN, input_layer):
  forward_propagation(NN, input_layer)
  return NN['layers'][-1]

def evaluate(NN, tags):
  output_layer = NN['layers'][-1]
  loss = np.mean(0.5*np.linalg.norm(output_layer - tags, axis=1)**2)

  y_pred = np.argmax(output_layer, axis=1)
  y_true = np.argmax(tags, axis=1)
  acc = accuracy_score(y_true, y_pred)

  # print(f'Loss = {loss}, acc = {acc}.')
  return (loss, acc)

if __name__ == '__main__':
  from sklearn import datasets, linear_model
  import matplotlib.pyplot as plt
  # generate sample data
  np.random.seed(0)
  X, y_true = datasets.make_moons(400, noise=0.20)
  tags = np.zeros((X.shape[0], 2))
  tags[np.where(y_true == 0), 0] = 1
  tags[np.where(y_true == 1), 1] = 1

  x_train      = X[:200]
  y_train_tags = tags[:200]
  x_test       = X[200:]
  y_test_tags  = tags[200:]
  # plt.scatter(X[:, 0], X[:, 1], c=y, cmap=plt.cm.Spectral)
  # plt.show()

  # initialize NN model
  nodes_size = [2, 8, 2]
  NN = initialize(nodes_size)
  # train
  y_res = train(NN, x_train, y_train_tags,
                eps=0.1, iter_num=100000,
                eval_num=1000, batch_size=-1)
  y_train_pred = np.argmax(y_res, axis=1)
  (loss, acc)  = evaluate(NN, y_train_tags)
  print('--------------------------------')
  print(f'train: Loss={loss}, acc={acc}.')

  # predict result
  y_res       = predict(NN, x_test)
  y_test_pred = np.argmax(y_res, axis=1)

  # plot
  plt.scatter(x_test[:, 0], x_test[:, 1], c=y_true[200:], cmap=plt.cm.Spectral)
  plt.title("ground truth")
  plt.show()

  plt.scatter(x_test[:, 0], x_test[:, 1], c=y_test_pred, cmap=plt.cm.Spectral)
  plt.title("predicted")
  plt.show()

  (loss, acc) = evaluate(NN, y_test_tags)
  print('--------------------------------')
  print(f'predict: Loss={loss}, acc={acc}.')