#-------------------------------------------------------------
#
# 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.
#
#-------------------------------------------------------------
args <- commandArgs(TRUE)
library("Matrix")
library("matrixStats") 

pad_image <- function(img, Hin, Win, padh, padw){
  C = nrow(img)
  img_padded = matrix(0, C, (Hin+2*padh)*(Win+2*padw), byrow=TRUE)  # zeros
  for (c in 1:C) {
    img_slice = matrix(img[c,], Hin, Win, byrow=TRUE)  # depth slice C reshaped
    img_padded_slice = matrix(0, Hin+2*padh, Win+2*padw)
    img_padded_slice[(padh+1):(padh+Hin), (padw+1):(padw+Win)] = img_slice
    img_padded[c,] = matrix(t(img_padded_slice), 1, (Hin+2*padh)*(Win+2*padw))  # reshape
  }
  img_padded
}

im2col <- function(img, Hin, Win, Hf, Wf, strideh, stridew) {
  C = nrow(img)
  Hout = as.integer((Hin - Hf) / strideh + 1)
  Wout = as.integer((Win - Wf) / stridew + 1)

  img_cols = matrix(0, C*Hf*Wf, Hout*Wout, byrow=TRUE)  # zeros
  for (hout in 1:Hout) {  # all output rows
    hin = (hout-1) * strideh + 1
    for (wout in 1:Wout) {  # all output columns
      win = (wout-1) * stridew + 1
      # Extract a local patch of the input image corresponding spatially to the filter sizes.
      img_patch = matrix(0, C, Hf*Wf, byrow=TRUE)  # zeros
      for (c in 1:C) {  # all channels
        img_slice = matrix(img[c,], Hin, Win, byrow=TRUE)  # reshape
        img_patch[c,] = matrix(t(img_slice[hin:(hin+Hf-1), win:(win+Wf-1)]), 1, Hf*Wf)
      }
      img_cols[,(hout-1)*Wout + wout] = matrix(t(img_patch), C*Hf*Wf, 1)  # reshape
    }
  }
  img_cols
}
		
conv2d <- function(X, W, C, Hin, Win, Hf, Wf, strideh, stridew, padh, padw) {
  N = nrow(X)
  F = nrow(W)
  Hout = as.integer((Hin + 2 * padh - Hf) / strideh + 1)
  Wout = as.integer((Win + 2 * padw - Wf) / stridew + 1)
  
  # Create output volume
  out = matrix(0, N, F*Hout*Wout, byrow=TRUE)

  # Convolution - im2col implementation
  for (n in 1:N) {  # all examples
    Xn = matrix(X[n,], C, Hin*Win, byrow=TRUE)  # reshape

    # Pad image
    Xn_padded = pad_image(Xn, Hin, Win, padh, padw)  # shape (C, (Hin+2*padh)*(Win+2*padw))

    # Extract local image patches into columns with im2col, of shape (C*Hf*Wf, Hout*Wout)
    Xn_padded_cols = im2col(Xn_padded, Hin+2*padh, Win+2*padw, Hf, Wf, strideh, stridew)

    # Convolve patches with filters
    outn = W %*% Xn_padded_cols   # shape (F, Hout*Wout)
    out[n,] = matrix(t(outn), 1, F*Hout*Wout)  # reshape
  }
  
  out
}

imgSize=8
numImg=16
numChannels=4
numFilters=3
filterSize=4
stride=1
pad=0

Hout = as.integer((imgSize + 2 * pad - filterSize) / stride + 1)

X = matrix(seq(1, numImg*numChannels*imgSize*imgSize), numImg, numChannels*imgSize*imgSize, byrow=TRUE);
W1 = matrix(seq(1, numFilters*numChannels*filterSize*filterSize), numFilters, numChannels*filterSize*filterSize, byrow=TRUE)
W2 = matrix(seq(1, numFilters*numChannels*filterSize*filterSize)+7, numFilters, numChannels*filterSize*filterSize, byrow=TRUE)
b = matrix(seq(1, numFilters), numFilters, 1, byrow=TRUE) 

X = X - rowMeans(X)

R1 = conv2d(X, W1, numChannels, imgSize, imgSize, filterSize, filterSize, stride, stride, pad, pad);
R2 = conv2d(X, W2, numChannels, imgSize, imgSize, filterSize, filterSize, stride, stride, pad, pad);
R = R1 + R2;

writeMM(as(R,"CsparseMatrix"), paste(args[2], "S", sep=""))