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

# The Linearized Image Transform function applies an affine transformation to linearized images.
# Optionally resizes the image (without scaling).
# Uses nearest neighbor sampling.
#
# .. code-block:: python
#
#   >>> import numpy as np
#   >>> from systemds.context import SystemDSContext
#   >>> from systemds.operator.algorithm import img_transform_linearized
#   >>> 
#   >>> with SystemDSContext() as sds:
#   ...     img = sds.from_numpy(
#   ...         np.array([[ 10., 20., 30.,
#   ...                     40., 50., 60.,
#   ...                     70., 80., 90. ]], dtype=np.float32)
#   ...     )
#   ...     result_img = img_transform_linearized(img, 3, 3, -1., 0., 2., 0., 1., 0., 255., 3, 3).compute()
#   ...     print(result_img.reshape(3, 3))
#   [[ 20.  10. 255.]
#    [ 50.  40. 255.]
#    [ 80.  70. 255.]]
#
#
# INPUT:
# -------------------------------------------------------------------------------------------
# img_in       Input images as linearized 2D matrix with top left corner at [1, 1] (every row represents a linearized matrix/image)
# out_w        Width of the output matrix
# out_h        Height of the output matrix
# a,b,c,d,e,f  The first two rows of the affine matrix in row-major order
# fill_value   The background of an image
# s_cols       Width of a single image
# s_rows       Height of a single image
# -------------------------------------------------------------------------------------------
#
# OUTPUT:
# ---------------------------------------------------------------------------------------
# img_out  Output images in linearized form as 2D matrix with top left corner at [1, 1]
# ---------------------------------------------------------------------------------------

m_img_transform_linearized = function(Matrix[Double] img_in, Integer out_w, Integer out_h, Double a, Double b, Double c, Double d,
 Double e, Double f, Double fill_value, Integer s_cols, Integer s_rows) return (Matrix[Double] img_out) {
    # size of a single image is s_cols : s_rows
  divisor = a * e - b * d
  if(divisor == 0) {
    print("Inverse matrix does not exist! Returning input.")
    img_out = img_in
  }
  else {
    orig_w = s_cols
    orig_h = s_rows
    # inverted transformation matrix
    # inversion is necessary because we compute the sampling position of pixels in the output image
    # and not the output coordinates of input pixels
    T_inv = matrix(0, rows=3, cols=3)
    T_inv[1, 1] = e / divisor
    T_inv[1, 2] = -b / divisor
    T_inv[1, 3] = (b * f - c * e) / divisor
    T_inv[2, 1] = -d / divisor
    T_inv[2, 2] = a / divisor
    T_inv[2, 3] = (c * d - a * f) / divisor
    T_inv[3, 3] = 1

    # coordinates of output pixel-centers linearized in row-major order
    coords = matrix(1, rows=3, cols=out_w*out_h)
    coords[1,] = t((seq(0, out_w*out_h-1) %% out_w) + 0.5)
    coords[2,] = t((seq(0, out_w*out_h-1) %/% out_w) + 0.5)

    # compute sampling pixel indices
    coords = floor(T_inv %*% coords) + 1

    inx = t(coords[1,])
    iny = t(coords[2,])

    # any out-of-range pixels, if present, correspond to an extra pixel with fill_value at the end of the input
    index_vector = (orig_w *(iny-1) + inx) * ((0<inx) & (inx<=orig_w) & (0<iny) & (iny<=orig_h)) 
    index_vector = t(index_vector)
    xs = ((index_vector == 0)*(orig_w*orig_h +1)) + index_vector
    
    if(min(index_vector) == 0){
      ys=cbind(img_in, matrix(fill_value,nrow(img_in), 1))
    }else{
      ys = img_in
    }

    ind= matrix(seq(1,ncol(xs),1),1,ncol(xs))
    z = table(xs, ind) 
    output = ys%*%z

    img_out = matrix(output, rows=nrow(img_in), cols=out_w*out_h)
  }
}
