#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
#-------------------------------------------------------------

# Trains a 2-layer autoencoder with minibatch SGD and step-size decay.
# If invoked with H1 > H2 then it becomes a 'bowtie' structured autoencoder
# Weights are initialized using Glorot & Bengio (2010) AISTATS initialization.
# The script standardizes the input before training (can be turned off).
# Also, it randomly reshuffles rows before training.
# Currently, tanh is set to be the activation function. 
# By re-implementing 'func' DML-bodied function, one can change the activation.
#
# INPUT:
# ---------------------------------------------------------------------------------------------
# X             Filename where the input is stored
# num_hidden1   Number of neurons in the 1st hidden layer
# num_hidden2   Number of neurons in the 2nd hidden layer
# max_epochs    Number of epochs to train for
# full_obj      If TRUE, Computes objective function value (squared-loss)
#               at the end of each epoch. Note that, computing the full 
#               objective can take a lot of time. 
# batch_size    Mini-batch size (training parameter)
# step          Initial step size (training parameter)
# decay         Decays step size after each epoch (training parameter)
# mu            Momentum parameter (training parameter)
# W1_rand       Weights might be initialized via input matrices
# W2_rand       --- 
# W3_rand       ---
# W4_rand       ---
# ---------------------------------------------------------------------------------------------
#
# OUTPUT:
# ----------------------------------------------------------------------------------------------------
# W1_out   Matrix storing weights between input layer and 1st hidden layer
# b1_out   Matrix storing bias between input layer and 1st hidden layer
# W2_out   Matrix storing weights between 1st hidden layer and 2nd hidden layer
# b2_out   Matrix storing bias between 1st hidden layer and 2nd hidden layer
# W3_out   Matrix storing weights between 2nd hidden layer and 3rd hidden layer
# b3_out   Matrix storing bias between 2nd hidden layer and 3rd hidden layer
# W4_out   Matrix storing weights between 3rd hidden layer and output layer
# b4_out   Matrix storing bias between 3rd hidden layer and output layer
# HIDDEN   Matrix storing the hidden (2nd) layer representation if needed
# ----------------------------------------------------------------------------------------------------

m_autoencoder_2layer = function(Matrix[Double] X, Integer num_hidden1, Integer num_hidden2, Integer max_epochs,
  Boolean full_obj = FALSE, Integer batch_size = 256, Double step = 1e-5, Double decay = 0.95, Double mu = 0.9,
  Matrix[Double] W1_rand = matrix(0, rows=0, cols=0), Matrix[Double] W2_rand = matrix(0, rows=0, cols=0), 
  Matrix[Double] W3_rand = matrix(0, rows=0, cols=0), Matrix[Double] W4_rand = matrix(0, rows=0, cols=0), 
  Matrix[Double] order_rand = matrix(0, rows=0, cols=0))
  return(Matrix[Double] W1, Matrix[Double] b1, Matrix[Double] W2, Matrix[Double] b2,
    Matrix[Double] W3, Matrix[Double] b3, Matrix[Double] W4, Matrix[Double] b4, Matrix[Double] reordered_H)
{
  n = nrow(X)
  m = ncol(X)
  #randomly reordering rows
  if(nrow(order_rand) == 0 & ncol(order_rand) == 0) 
    permut = table(seq(1,n,1), order(target=Rand(rows=n, cols=1, min=0, max=1, pdf="uniform"), by=1, index.return=TRUE), n, n)
  else 
    permut = table(seq(1,n,1), order(target=order_rand, by=1, index.return=TRUE), n, n)
  X = permut %*% X

  #z-transform, whitening operator is better
  means = colSums(X)/n
  stds = sqrt((colSums(X^2)/n - means*means)*n/(n-1)) + 1e-17
  X = (X - means)/stds
  
  if(nrow(W1_rand) == 0 & ncol(W1_rand) == 0) 
    W1_rand = Rand(rows=num_hidden1, cols=m, min=-1, max=1, pdf="uniform")
  if(nrow(W2_rand) == 0 & ncol(W2_rand) == 0) 
    W2_rand = Rand(rows=num_hidden2, cols=num_hidden1, min=-1, max=1, pdf="uniform")
  if(nrow(W3_rand) == 0 & ncol(W3_rand) == 0) 
    W3_rand = Rand(rows=num_hidden1, cols=num_hidden2, min=-1, max=1, pdf="uniform")
  if(nrow(W4_rand) == 0 & ncol(W4_rand) == 0) 
    W4_rand = Rand(rows=m, cols=num_hidden1, min=-1, max=1, pdf="uniform")

  W1 = sqrt(6)/sqrt(m + num_hidden1) * W1_rand 
  b1 = matrix(0, rows=num_hidden1, cols=1)
  W2 = sqrt(6)/sqrt(num_hidden1 + num_hidden2) * W2_rand 
  b2 = matrix(0, rows=num_hidden2, cols=1)
  W3 = sqrt(6)/sqrt(num_hidden2 + num_hidden1) * W3_rand 
  b3 = matrix(0, rows=num_hidden1, cols=1)
  W4 = sqrt(6)/sqrt(num_hidden2 + m) * W4_rand 
  b4 = matrix(0, rows=m, cols=1)

  upd_W1 = matrix(0, rows=nrow(W1), cols=ncol(W1))
  upd_b1 = matrix(0, rows=nrow(b1), cols=ncol(b1))
  upd_W2 = matrix(0, rows=nrow(W2), cols=ncol(W2))
  upd_b2 = matrix(0, rows=nrow(b2), cols=ncol(b2))
  upd_W3 = matrix(0, rows=nrow(W3), cols=ncol(W3))
  upd_b3 = matrix(0, rows=nrow(b3), cols=ncol(b3))
  upd_W4 = matrix(0, rows=nrow(W4), cols=ncol(W4))
  upd_b4 = matrix(0, rows=nrow(b4), cols=ncol(b4))

  if( full_obj ){
    [full_H1, full_H1_prime, full_H2, full_H2_prime, full_H3, full_H3_prime, full_Yhat, 
     full_Yhat_prime, full_E] = feedForward(X, W1, b1, W2, b2, W3, b3, W4, b4, X)
    full_o = obj(full_E)
    print("EPOCHS=" + 0 + " OBJ (FULL DATA): " + full_o)
  }

  iter = 0
  num_iters_per_epoch = ceil(n / batch_size)
  max_iterations = max_epochs * num_iters_per_epoch
  #print("num_iters_per_epoch=" + num_iters_per_epoch + " max_iterations=" + max_iterations)
  beg = 1
  while( iter < max_iterations ){
    end = beg + batch_size - 1
    if(end > n) end = n
    X_batch = X[beg:end,]

    [H1, H1_prime, H2, H2_prime, H3, H3_prime, Yhat, Yhat_prime, E] 
      = feedForward(X_batch, W1, b1, W2, b2, W3, b3, W4, b4, X_batch)
    [W1_grad, b1_grad, W2_grad, b2_grad, W3_grad, b3_grad, W4_grad, b4_grad] 
      = grad(X_batch, H1, H1_prime, H2, H2_prime, H3, H3_prime, Yhat_prime, E, W1, W2, W3, W4)

    o = obj(E)
    print("epochs=%5.4f BATCH beg=%d end=%d obj=%f", (iter / num_iters_per_epoch), beg, end, o)

    #update
    local_step = step / nrow(X_batch)
    upd_W1 = mu * upd_W1 - local_step * W1_grad
    upd_b1 = mu * upd_b1 - local_step * b1
    upd_W2 = mu * upd_W2 - local_step * W2_grad
    upd_b2 = mu * upd_b2 - local_step * b2
    upd_W3 = mu * upd_W3 - local_step * W3_grad
    upd_b3 = mu * upd_b3 - local_step * b3
    upd_W4 = mu * upd_W4 - local_step * W4_grad
    upd_b4 = mu * upd_b4 - local_step * b4
    W1 = W1 + upd_W1
    b1 = b1 + upd_b1
    W2 = W2 + upd_W2
    b2 = b2 + upd_b2
    W3 = W3 + upd_W3
    b3 = b3 + upd_b3
    W4 = W4 + upd_W4
    b4 = b4 + upd_b4

    iter = iter + 1
    if(end == n) beg = 1
    else beg = end + 1

    if( iter %% num_iters_per_epoch == 0 ) step = step * decay

    if( full_obj & iter %% num_iters_per_epoch == 0 ){
      [full_H1, full_H1_prime, full_H2, full_H2_prime, full_H3, full_H3_prime, full_Yhat, 
       full_Yhat_prime, full_E] = feedForward(X, W1, b1, W2, b2, W3, b3, W4, b4, X)
      full_o = obj(full_E)
      epochs = iter %/% num_iters_per_epoch
      print("EPOCHS=" + epochs + " iter=" + iter + " OBJ (FULL DATA)=" + full_o)
    }
  }

  [full_H1, full_H1_prime, full_H2, full_H2_prime, full_H3, full_H3_prime, full_Yhat, 
   full_Yhat_prime, full_E] = feedForward(X, W1, b1, W2, b2, W3, b3, W4, b4, X)
  reordered_H = t(permut) %*% full_H2
}

#implements tanh
#to use another activation fn, implement a DML-bodied function with 
#function name 'func' and comment out this one
func = function(Matrix[Double] X) return(Matrix[Double] Y, Matrix[Double] Y_prime){
  Y = (exp(2*X) - 1)/(exp(2*X) + 1)
  Y_prime = 1 - Y^2
}

feedForward = function(Matrix[Double] X, Matrix[Double] W1, Matrix[Double] b1, Matrix[Double] W2, Matrix[Double] b2, 
  Matrix[Double] W3, Matrix[Double] b3, Matrix[Double] W4, Matrix[Double] b4, Matrix[Double] Y)
  return(Matrix[Double] H1, Matrix[Double] H1_prime,Matrix[Double] H2, Matrix[Double] H2_prime,
    Matrix[Double] H3, Matrix[Double] H3_prime, Matrix[Double] Yhat, Matrix[Double] Yhat_prime, Matrix[Double] E)
{
  H1_in = t(W1 %*% t(X) + b1)
  [H1, H1_prime] = func(H1_in)

  H2_in = t(W2 %*% t(H1) + b2)
  [H2, H2_prime] = func(H2_in)

  H3_in = t(W3 %*% t(H2) + b3)
  [H3, H3_prime] = func(H3_in)

  Yhat_in = t(W4 %*% t(H3) + b4)
  [Yhat, Yhat_prime] = func(Yhat_in)
  E = Yhat - Y
}

grad = function(Matrix[Double] X, Matrix[Double] H1, Matrix[Double] H1_prime, Matrix[Double] H2, Matrix[Double] H2_prime,
  Matrix[Double] H3, Matrix[Double] H3_prime, Matrix[Double] Yhat_prime, Matrix[Double] E,
  Matrix[Double] W1, Matrix[Double] W2, Matrix[Double] W3, Matrix[Double] W4)
  return(Matrix[Double] W1_grad, Matrix[Double] b1_grad,Matrix[Double] W2_grad, Matrix[Double] b2_grad, 
    Matrix[Double] W3_grad, Matrix[Double] b3_grad, Matrix[Double] W4_grad, Matrix[Double] b4_grad)
{
  #backprop
  delta4 = E * Yhat_prime
  delta3 = H3_prime * (delta4 %*% W4)
  delta2 = H2_prime * (delta3 %*% W3)
  delta1 = H1_prime * (delta2 %*% W2)

  #compute gradients
  b4_grad = t(colSums(delta4))
  b3_grad = t(colSums(delta3))
  b2_grad = t(colSums(delta2))
  b1_grad = t(colSums(delta1))

  W4_grad = t(delta4) %*% H3
  W3_grad = t(delta3) %*% H2
  W2_grad = t(delta2) %*% H1
  W1_grad = t(delta1) %*% X
}

obj = function(Matrix[Double] E) return(Double val){
  val = 0.5 * sum(E^2)
}
