#-------------------------------------------------------------
#
# 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 script performs forward pass through ff neural network in ffTrain and ffPredict builtins
#
# INPUT PARAMETERS:
# --------------------------------------------------------------------------------------------
# NAME      TYPE            DEFAULT     MEANING
# --------------------------------------------------------------------------------------------
# X         Matrix[double]    ---       Input data matrix, of shape (N, C*Hin*Win)
# C         Integer           ---       Number of input channels
# Hin       Integer           ---       Input height
# Win       Integer           ---       Input width
# Model     list[unknown]     ---       List of layers
# predict   Boolean           FALSE     Flag used to avoid dropout when predicting
# --------------------------------------------------------------------------------------------
# OUTPUT: 
# cache     List[unknown]     ---       Caches all intermediate steps of forward pass
#

# TODO remove from layers and move into FFN once sideway imports are working

source("nn/layers/affine.dml") as affine
source("nn/layers/conv2d_builtin.dml") as conv2d
source("nn/layers/dropout.dml") as dropout
source("nn/layers/max_pool2d_builtin.dml") as max_pool2d
source("nn/layers/relu.dml") as relu
source("nn/layers/softmax.dml") as softmax

lenetForward = function(Matrix[Double] X, Integer C,
  Integer Hin, Integer Win, list[unknown] model, Boolean predict = FALSE) 
  return (list[unknown] cache)
{
  Hf = 5  # filter height
  Wf = 5  # filter width
  stride = 1
  pad = 2  # For same dimensions, (Hf - stride) / 2

  f1 = as.matrix(model["W1"])
  f2 = as.matrix(model["W2"])
  F1 = nrow(f1)  # num conv filters in conv1
  F2 = nrow(f2)  # num conv filters in conv2

  # Compute forward pass
  ## layer 1: conv1 -> relu1 -> pool1
  [outc1, Houtc1, Woutc1] = conv2d::forward(X, as.matrix(model["W1"]), as.matrix(model["b1"]), 
                                                C, Hin, Win, Hf, Wf, stride, stride, pad, pad)
  outr1 = relu::forward(outc1)
  [outp1, Houtp1, Woutp1] = max_pool2d::forward(outr1, F1, Houtc1, Woutc1, Hf=2, Wf=2,
                                              strideh=2, stridew=2, padh=0, padw=0)
  ## layer 2: conv2 -> relu2 -> pool2
  [outc2, Houtc2, Woutc2] = conv2d::forward(outp1, as.matrix(model["W2"]), as.matrix(model["b2"]), 
                                            F1, Houtp1, Woutp1, Hf, Wf, stride, stride, pad, pad)
  outr2 = relu::forward(outc2)
  [outp2, Houtp2, Woutp2] = max_pool2d::forward(outr2, F2, Houtc2, Woutc2, Hf=2, Wf=2,
                                                strideh=2, stridew=2, padh=0, padw=0)
  ## layer 3:  affine3 -> relu3 -> dropout
  outa3 = affine::forward(outp2, as.matrix(model["W3"]), as.matrix(model["b3"]))
  outr3 = relu::forward(outa3)
  [outd3, maskd3] = dropout::forward(outr3, 0.5, -1)
  if(predict)
    outd3 = outr3
  ## layer 4:  affine4 -> softmax
  outa4 = affine::forward(outd3, as.matrix(model["W4"]), as.matrix(model["b4"]))
  probs = softmax::forward(outa4)

  cache = list(outc1=outc1, Houtc1=Houtc1, Woutc1=Woutc1, outr1=outr1, outp1=outp1, Houtp1=Houtp1, 
  Woutp1=Woutp1, outc2=outc2, Houtc2=Houtc2, Woutc2=Woutc2, outr2=outr2, outp2=outp2, Houtp2=Houtp2, 
  Woutp2=Woutp2, outa3=outa3, outr3=outr3, outd3=outd3, maskd3=maskd3, outa4=outa4, probs=probs)
}
