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

/*
 * Upsampling layer for 2D inputs.
 *
 * Repeats the rows and columns of the data by size_h and size_w respectively.
 */

forward = function(matrix[double] X, int C, int Hin, int Win, int size_h, int size_w)
    return (matrix[double] out) {
  /*
   * Computes the forward pass for a Upsampling layer.
   *
   *
   * Inputs:
   *  - X: Inputs, of shape (N, C*Hin*Win).
   *  - C: Number of input channels (dimensionality of input depth).
   *  - Hin: Input height.
   *  - Win: Input width.
   *  - size_h: upsampling factor for rows.
   *  - size_w: upsampling factor for columns.
   *
   * Outputs:
   *  - out: Outputs, of shape (N, C*Hout*Wout), where Hout = Hin*size_h, and Wout = Win * size_w.
   */
  N = nrow(X)
  Hout = size_h*Hin
  Wout = size_w*Win
  emptyInput = matrix(0, rows=N, cols=C*Hout*Wout)
  out = avg_pool_backward(emptyInput, X, input_shape=[N,C,Hout,Wout], pool_size=[size_h,size_w], stride=[size_h,size_w], padding=[0,0])
  out = out * size_h * size_w
}

backward = function(matrix[double] dout, int C, int Hin, int Win, int size_h, int size_w)
    return (matrix[double] dX) {
  /*
   * Computes the backward pass for a Upsampling layer.
   *
   * Inputs:
   *  - dout: Gradient wrt `out` from upstream.
   *  - C: Number of input channels (dimensionality of input depth).
   *  - Hin: Input height.
   *  - Win: Input width.
   *  - size_h: upsampling factor for rows.
   *  - size_w: upsampling factor for columns.
   *
   * Outputs:
   *  - dX: Gradient wrt `X`, of same shape as `X`.
   */
   N = nrow(dout)
   Hout = size_h*Hin
   Wout = size_w*Win
   dX = avg_pool(dout, input_shape=[N,C,Hout,Wout], pool_size=[size_h,size_w], stride=[size_h,size_w], padding=[0,0])
   dX = dX * size_h * size_w
}