#-------------------------------------------------------------
#
# 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.
#
#-------------------------------------------------------------

# This builtin function trains LeNet CNN. The architecture of the
# networks is:conv1 -> relu1 -> pool1 -> conv2 -> relu2 -> pool2 ->
# affine3 -> relu3 -> affine4 -> softmax
#
# INPUT:
# ----------------------------------------------------------
# X           Input data matrix, of shape (N, C*Hin*Win)
# Y           Target matrix, of shape (N, K)
# X_val       Validation data matrix, of shape (N, C*Hin*Win)
# Y_val       Validation target matrix, of shape (N, K)
# C           Number of input channels (dimensionality of input depth)
# Hin         Input width
# Win         Input height
# batch_size  Batch size
# epochs      Number of epochs
# lr          Learning rate
# mu          Momentum value
# decay       Learning rate decay
# reg         Regularization strength
# seed        Seed for model initialization
# verbose     Flag indicates if function should print to stdout
# ------------------------------------------------------------------------------------
#
# OUTPUT:
# -------------------------------------------------------------------------------
# model  Trained model which can be used in lenetPredict
# -------------------------------------------------------------------------------

source("nn/layers/affine.dml") as affine
source("nn/layers/conv2d_builtin.dml") as conv2d
source("nn/layers/cross_entropy_loss.dml") as cross_entropy_loss
source("nn/layers/dropout.dml") as dropout
source("nn/layers/l2_reg.dml") as l2_reg
source("nn/layers/max_pool2d_builtin.dml") as max_pool2d
source("nn/layers/relu.dml") as relu
source("nn/layers/softmax.dml") as softmax
source("nn/optim/sgd_nesterov.dml") as sgd_nesterov
source("nn/layers/lenetForwardPass.dml") as lenet_fw

m_lenetTrain = function(Matrix[Double] X, Matrix[Double] Y, Matrix[Double] X_val, 
  Matrix[Double] Y_val, Integer C, Integer Hin, Integer Win, Integer batch_size=64, 
  Integer epochs=20, Double lr=0.01, Double mu=0.9, Double decay=0.95, Double reg=5e-04, 
  Boolean verbose=FALSE, Integer seed=-1)
  return (List[unknown] model)
{

  N = nrow(X)
  K = ncol(Y)

  # Create network:
  # conv1 -> relu1 -> pool1 -> conv2 -> relu2 -> pool2 -> affine3 -> relu3 -> affine4 -> softmax
  Hf = 5  # filter height
  Wf = 5  # filter width
  stride = 1
  pad = 2  # For same dimensions, (Hf - stride) / 2

  F1 = 32  # num conv filters in conv1
  F2 = 64  # num conv filters in conv2
  N3 = 512  # num nodes in affine3
  # Note: affine4 has K nodes, which is equal to the number of target dimensions (num classes)

  [W1, b1] = conv2d::init(F1, C, Hf, Wf, seed)  # inputs: (N, C*Hin*Win)
  [W2, b2] = conv2d::init(F2, F1, Hf, Wf, seed)  # inputs: (N, F1*(Hin/2)*(Win/2))
  [W3, b3] = affine::init(F2*(Hin/2/2)*(Win/2/2), N3, seed)  # inputs: (N, F2*(Hin/2/2)*(Win/2/2))
  [W4, b4] = affine::init(N3, K, seed)  # inputs: (N, N3)
  W4 = W4 / sqrt(2)  # different initialization, since being fed into softmax, instead of relu

  # Initialize SGD w/ Nesterov momentum optimizer
  
  vW1 = sgd_nesterov::init(W1); vb1 = sgd_nesterov::init(b1)
  vW2 = sgd_nesterov::init(W2); vb2 = sgd_nesterov::init(b2)
  vW3 = sgd_nesterov::init(W3); vb3 = sgd_nesterov::init(b3)
  vW4 = sgd_nesterov::init(W4); vb4 = sgd_nesterov::init(b4)

  # Optimize  
  iters = ceil(N / batch_size)
  for (e in 1:epochs) {
    loss = 0
    val_loss = 0
    val_acc = 0
    for(i in 1:iters) {
      # Get next batch
      beg = ((i-1) * batch_size) %% N + 1
      end = min(N, beg + batch_size - 1)
      X_batch = X[beg:end,]
      y_batch = Y[beg:end,]

      model = list(W1=W1, b1=b1, W2=W2, b2=b2, W3=W3, b3=b3, W4=W4, b4=b4)

      cache = lenet_fw::lenetForward(X=X_batch, C=C, Hin=Hin, Win=Win, model=model)
      
      # loss
      if(verbose){
        loss = loss + cross_entropy_loss::forward(as.matrix(cache["probs"]), y_batch)
        val_cache = lenet_fw::lenetForward(X=X_val, C=C, Hin=Hin, Win=Win, model=model)
        val_loss = val_loss + cross_entropy_loss::forward(as.matrix(val_cache["probs"]), Y_val)
        val_acc = val_acc + mean(rowIndexMax(as.matrix(val_cache["probs"])) == rowIndexMax(Y_val))
      }
      
      dprobs = cross_entropy_loss::backward(as.matrix(cache["probs"]), y_batch)

      # Compute data backward pass
      [dW1, db1, dW2, db2, dW3, db3, dW4, db4] = feed_backward(
        X_batch, C, Hin, Win, reg, model, dprobs, cache)

      # Optimize with SGD w/ Nesterov momentum
      [W1, vW1] = sgd_nesterov::update(W1, dW1, lr, mu, vW1)
      [b1, vb1] = sgd_nesterov::update(b1, db1, lr, mu, vb1)
      [W2, vW2] = sgd_nesterov::update(W2, dW2, lr, mu, vW2)
      [b2, vb2] = sgd_nesterov::update(b2, db2, lr, mu, vb2)
      [W3, vW3] = sgd_nesterov::update(W3, dW3, lr, mu, vW3)
      [b3, vb3] = sgd_nesterov::update(b3, db3, lr, mu, vb3)
      [W4, vW4] = sgd_nesterov::update(W4, dW4, lr, mu, vW4)
      [b4, vb4] = sgd_nesterov::update(b4, db4, lr, mu, vb4)
    }
    
    if(verbose) {
      print("Epoch: " + e + ", Train loss: " + loss/iters + ", Validation loss: " + 
      val_loss/iters + ", Validation accuracy: " + val_acc/iters)
    }
    # Anneal momentum towards 0.999
    mu = mu + (0.999 - mu)/(1+epochs-e)
    # Decay learning rate
    lr = lr * decay
  }

  model = list(W1=W1, b1=b1, W2=W2, b2=b2, W3=W3, b3=b3, W4=W4, b4=b4)
}

feed_backward = function(Matrix[Double] X, Integer C, Integer Hin, Integer Win, 
  Double reg,list[unknown] model, matrix[Double] dprobs, list[unknown] cache)
  return (Matrix[Double] dW1, Matrix[Double] db1,
  Matrix[Double] dW2, Matrix[Double] db2,
  Matrix[Double] dW3, Matrix[Double] db3,
  Matrix[Double] dW4, Matrix[Double] db4)
{
  Hf = 5  # filter height
  Wf = 5  # filter width
  stride = 1
  pad = 2  # For same dimensions, (Hf - stride) / 2

  F1 = 32  # num conv filters in conv1
  F2 = 64  # num conv filters in conv2

  # Compute data backward pass

  ## layer 4:  affine4 -> softmax
  douta4 = softmax::backward(dprobs, as.matrix(cache["outa4"]))
  [doutd3, dW4, db4] = affine::backward(douta4, as.matrix(cache["outd3"]), as.matrix(model["W4"]), as.matrix(model["b4"]))
  ## layer 3:  affine3 -> relu3 -> dropout
  doutr3 = dropout::backward(doutd3, as.matrix(cache["outr3"]), 0.5, as.matrix(cache["maskd3"]))
  douta3 = relu::backward(doutr3, as.matrix(cache["outa3"]))
  [doutp2, dW3, db3] = affine::backward(douta3, as.matrix(cache["outp2"]), as.matrix(model["W3"]), as.matrix(model["b3"]))
  ## layer 2: conv2 -> relu2 -> pool2
  doutr2 = max_pool2d::backward(doutp2, as.integer(as.scalar(cache["Houtp2"])), as.integer(as.scalar(cache["Woutp2"])), 
  as.matrix(cache["outr2"]), F2, as.integer(as.scalar(cache["Houtc2"])), as.integer(as.scalar(cache["Woutc2"])), Hf=2, Wf=2,
                                                                                        strideh=2, stridew=2, padh=0, padw=0)
  doutc2 = relu::backward(doutr2, as.matrix(cache["outc2"]))
  [doutp1, dW2, db2] = conv2d::backward(doutc2, as.integer(as.scalar(cache["Houtc2"])), as.integer(as.scalar(cache["Woutc2"])), 
  as.matrix(cache["outp1"]), as.matrix(model["W2"]), as.matrix(model["b2"]), F1, as.integer(as.scalar(cache["Houtp1"])), 
                                              as.integer(as.scalar(cache["Woutp1"])), Hf, Wf, stride, stride, pad, pad)
  ## layer 1: conv1 -> relu1 -> pool1
  doutr1 = max_pool2d::backward(doutp1, as.integer(as.scalar(cache["Houtp1"])), as.integer(as.scalar(cache["Woutp1"])), 
  as.matrix(cache["outr1"]), F1, as.integer(as.scalar(cache["Houtc1"])), as.integer(as.scalar(cache["Woutc1"])), Hf=2, Wf=2,
                                                                                      strideh=2, stridew=2, padh=0, padw=0)
  doutc1 = relu::backward(doutr1, as.matrix(cache["outc1"]))
  [dX_batch, dW1, db1] = conv2d::backward(doutc1, as.integer(as.scalar(cache["Houtc1"])), as.integer(as.scalar(cache["Woutc1"])), 
                                X, as.matrix(model["W1"]), as.matrix(model["b1"]), C, Hin, Win, Hf, Wf, stride, stride, pad, pad)

  # Compute regularization backward pass
  dW1_reg = l2_reg::backward(as.matrix(model["W1"]), reg)
  dW2_reg = l2_reg::backward(as.matrix(model["W2"]), reg)
  dW3_reg = l2_reg::backward(as.matrix(model["W3"]), reg)
  dW4_reg = l2_reg::backward(as.matrix(model["W4"]), reg)
  dW1 = dW1 + dW1_reg
  dW2 = dW2 + dW2_reg
  dW3 = dW3 + dW3_reg
  dW4 = dW4 + dW4_reg
}
