#-------------------------------------------------------------
#
# 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 raGroupby-function takes a matrix dataset as input from where it performs
# relational operations : groupby
#
# INPUT:
# ------------------------------------------------------------------------------
# X         Matrix of input data [shape: N x M]
# col       Integer indicating the column index to execute grupby command
# method    Groupby implemention method (nested-loop, permutation-matrix)
# ------------------------------------------------------------------------------
#
# OUTPUT:
# ------------------------------------------------------------------------------
# Y         Matrix of selected data [shape N' x M] with N' <= N
# ------------------------------------------------------------------------------

m_raGroupby = function (Matrix[Double] X, Integer col, String method)
  return (Matrix[Double] Y)
{
    if (method == "nested-loop") {
    # Extract and sort unique group values from the specified column (1-based index)
    groupsUnique = unique(X[, col])
    groupsUniqueOrdered = order(target = groupsUnique, by = 1)
    numGroups = nrow(groupsUnique)
    maxRowsInGroup = max(table(X[,col],1));

    # Define a zero output matrix, save the initial order of the groups, and sort increasingly
    Y = matrix(0, numGroups, maxRowsInGroup*(ncol(X) - 1) + 1)
    Y[,1] = groupsUnique
    indicesY = order(target = Y, by = 1, index.return = TRUE)
    Y = order(target = Y, by = 1, decreasing = FALSE, index.return = FALSE)

    # Order the input matrix by the grouping column
    indicesX = order(target = X, by = col, index.return = TRUE)
    X = order(target = X, by = col, decreasing = FALSE, index.return = FALSE)

    currentGroupX = 1
    currentGroupY = 1
    i = 1

    # Iterate over the input matrix
    while (numGroups > 0) {
        currentGroup = as.scalar(Y[currentGroupX,1])
        nRowsToCopy = 0

        # Find the rows for the current group
        group = 1
        while (group > 0) {
            # Break if there are no more rows left in X
            if (i > nrow(X)) {
                group = 0
            }
            # Check if the row belongs to the current group
            else if (as.scalar(X[i, col]) == currentGroup) {
                nRowsToCopy = nRowsToCopy + 1
                i = i + 1
            }
            # Break if the row does not belong to the current group
            else {
               group = 0
            }
        }

        # Copy the values into the output matrix
        if (nRowsToCopy > 0) {
            nRowsCurrentGroup = currentGroupY + nRowsToCopy - 1

            # 1. Grouping column is the first column
            if (col == 1) {
                newMatrix = X[currentGroupY:nRowsCurrentGroup, (col+1):ncol(X)]
            }
            # 2. Grouping column is the last column
            else if (col == ncol(X)) {
                newMatrix = X [currentGroupY:nRowsCurrentGroup, 1:col-1]
            }
            # 3. Grouping column has an intermediate position
            else {
                newMatrix = cbind(X[currentGroupY:nRowsCurrentGroup, 1:(col-1)], X[currentGroupY:nRowsCurrentGroup, (col+1):ncol(X)])
            }

            # Flatten the new row
            newRow = matrix(newMatrix, rows = 1, cols = nrow(newMatrix) * ncol(newMatrix))
            newRowColIdx = nRowsToCopy * (ncol(X)-1)

            # Add the new row into Y at the current group
            Y[currentGroupX, 2: (newRowColIdx + 1)] = newRow
        }

        # Continue with the next group
        currentGroupX = currentGroupX + 1
        currentGroupY = currentGroupY + nRowsToCopy
        numGroups = numGroups - 1
    }

    # Restore the initial order of X
    X = cbind(X, indicesX)
    nColX = ncol(X)
    X = order(target = X, by= nColX)
    X = X[, 1:nColX-1]

    # Restore the initial order of Y
    Y = cbind(Y, indicesY)
    nColY = ncol(Y)
    Y = order(target = Y, by= nColY)
    Y = Y[, 1:nColY-1]
  }

  else if (method == "permutation-matrix") {
    # Extract the grouping column and create unique groups
    key = X[,col]
    keyUnique = unique(X[, col])
    numGroups = nrow(keyUnique)
    maxRowsInGroup = max(table(X[,col],1))

    # Calculate the frequency of each group
    freqPerKey = table(key, 1)
    freqPerKey = removeEmpty(target = freqPerKey, margin = "rows")
    freqPerKeyIndices = order(target = keyUnique, by = 1, index.return = TRUE)

    # Match the length of freqPerKey to keyUnique and sort it accordingly
    freqPerKey = cbind(freqPerKey, freqPerKeyIndices)
    nColFpk = ncol(freqPerKey)
    freqPerKey = order(target = freqPerKey, by= nColFpk)
    freqPerKey = freqPerKey[, 1:nColFpk-1]
    freqPerKey = t(freqPerKey)

    # Find the group with the most values
    groupMaxVal = maxRowsInGroup*(ncol(X)-1)+1
    groupMaxValKey = max(freqPerKey)

    # Calculate the amount of rows that need padding and the amount of padding per key
    groupMaxValKeySeq = matrix(groupMaxValKey, nrow(freqPerKey), ncol(freqPerKey))
    missingPadding = groupMaxValKeySeq - freqPerKey
    amountOfZeroRows = sum(missingPadding)

    # 1. Padding is required
    if (amountOfZeroRows > 0) {
        missingPadding = t(missingPadding)

        # Remove the keys that dont need padding
        removeMask = (missingPadding != 0)
        missingPadding = cbind(keyUnique, missingPadding)
        missingPadding = removeEmpty(target = missingPadding, margin = "rows", select = removeMask)

        # Keys that need padding and padding length per group
        keysPadding = missingPadding[,1]
        missingPadding = missingPadding[,2]
        repeatKeys = matrix(0, rows=amountOfZeroRows, cols=1)

        # Generate the repeating keys
        repeatKeysIdxS = 1

        for (i in 1:nrow(missingPadding)) {
           repeat_count = as.scalar(missingPadding[i,1])
           if (repeat_count > 0) {
              temp = matrix(as.scalar(keysPadding[i, 1]), rows=repeat_count, cols = 1)
              repeatKeysIdxE = repeatKeysIdxS + repeat_count - 1
              repeatKeys[repeatKeysIdxS:repeatKeysIdxE, 1] = temp
              repeatKeysIdxS = repeatKeysIdxE + 1
           }
        }

        # Combine the keys that need padding with the actual padding
        padding = matrix(0, rows = nrow(repeatKeys), cols = 1)
        padding = cbind(repeatKeys, padding)

        # Extend the existing keys to a second column to match the padded keys
        key = key %*% matrix(1, rows = 1, cols = 2)

        # Combine the keys with the padded keys and sort them increasingly
        tempY = rbind(key, padding)
        tempY = order(target = tempY, by = 1, decreasing = FALSE, index.return = FALSE)

        # Remove the padded rows and save the Indices of the combined keys for the permutation matrix
        paddedRows = tempY[, 2]
        tempIndicesY = order(target = tempY, by = 1, decreasing = FALSE, index.return = TRUE)
        tempIndicesY = removeEmpty(target = tempIndicesY, margin = "rows", select = (paddedRows!=0))

        # Create the permutation matrix by using the Indices of the combined keys
        P = table(seq(1, nrow(X)), tempIndicesY)

        # Order the initial matrix to match the sorted keys with padding
        indicesX = order(target = X, by = col, index.return = TRUE)
        X = order(target = X, by = col, decreasing = FALSE, index.return = FALSE)
        X = order(target = X, by = col, decreasing = FALSE, index.return = FALSE)

        # Perform the matrix multiplication
        tempY = t(P) %*% X
    }

    # 2. Padding is not required
    else {
        tempY = X
        tempY = order(target = tempY, by = col, decreasing = FALSE, index.return = FALSE)
    }

    # Remove the selected column from tempY
    if (col == 1) {
        tempY = tempY[, col+1:ncol(tempY)]
    }
    else if (col == ncol(X)) {
        tempY = tempY[, 1:col-1]
    }
    else {
        tempY = cbind(tempY[, 1:col-1],tempY[, col+1:ncol(tempY)])
    }

    # Set the value of the final output
    Y = matrix(0, rows=numGroups, cols=groupMaxVal)
    Y[,1] = keyUnique

    # Each group's data may not fill exactly maxRowsInGroup rows
    # If needed, we need to pad to the expected size first
    expectedRows = numGroups * maxRowsInGroup
    actualRows = nrow(tempY)

    if (actualRows < expectedRows) {
      # Pad tempY with zeros to match expected structure
      tempYPadded = matrix(0, rows=expectedRows, cols=ncol(tempY))
      tempYPadded[1:actualRows,] = tempY
    }
    else {
      tempYPadded = tempY
    }

    # Save the initial order of the groups in Y and order Y to match the sorted tempYPadded
    indicesY = order(target = Y, by = 1, index.return = TRUE)
    Y = order(target = Y, by = 1, decreasing = FALSE, index.return = FALSE)

    # Copy the values into Y
    Y[,2:ncol(Y)] = matrix(tempYPadded, rows=numGroups, cols=groupMaxVal-1)

    # Restore the initial order of X
    X = cbind(X, indicesX)
    nColX = ncol(X)
    X = order(target = X, by= nColX)
    X = X[, 1:nColX-1]

    # Restore the initial order of Y
    Y = cbind(Y, indicesY)
    nColY = ncol(Y)
    Y = order(target = Y, by= nColY)
    Y = Y[, 1:nColY-1]
  }
}
