#-------------------------------------------------------------
#
# 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 makes prediction given data and trained LeNet model
#
# INPUT:
# ---------------------------------------------------------------------------------------------
# model       Trained LeNet model
# X           Input data matrix, of shape (N, C*Hin*Win)
# C           Number of input channels
# Hin         Input height
# Win         Input width
# batch_size  Batch size
# ---------------------------------------------------------------------------------------------
#
# OUTPUT:
# ----------------------------------------------------------------------------------------
# probs  Predicted values
# ----------------------------------------------------------------------------------------

source("nn/layers/lenetForwardPass.dml") as lenet_fw

s_lenetPredict = function(list[unknown] model, Matrix[Double] X, Integer C,
  Integer Hin, Integer Win, Integer batch_size = 128)
  return (Matrix[Double] probs)
{
  N = nrow(X)

  # 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
  W1 = as.matrix(model["W1"])
  W2 = as.matrix(model["W2"])
  W3 = as.matrix(model["W3"])
  W4 = as.matrix(model["W4"])
  F1 = nrow(W1)  # num conv filters in conv1
  F2 = nrow(W2)  # num conv filters in conv2
  N3 = ncol(W3)  # num nodes in affine3
  K = ncol(W4)  # num nodes in affine4, equal to number of target dimensions (num classes)

  # Compute predictions over mini-batches
  probs = matrix(0, rows=N, cols=K)
  iters = ceil(N / batch_size)
  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,]

    pred = lenet_fw::lenetForward(X_batch, C, Hin, Win, model, TRUE)

    # Store predictions
    probs[beg:end,] = as.matrix(pred["probs"])
  }
}
