#-------------------------------------------------------------
#
# 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 SCRIPT COMPUTES A SOLUTION FOR THE STABLE MARRIAGE PROBLEM (WITHOUT TIES, COMPLETE LISTS)
#
# INPUT PARAMETERS:
# --------------------------------------------------------------------------------------------
# NAME                  TYPE    DEFAULT    MEANING
# --------------------------------------------------------------------------------------------
# P                     String     ---   Location (on HDFS) to read the proposer's preference matrix P.
#                                        P is assumed to store preferences row-wise (i.e. row j stores proposer j's preferences).
#                                        It must be a square matrix with no zeros.
#
# A                     String     ---   Location (on HDFS) to read the acceptor's preference matrix A.
#                                        A is assumed to store preferences row-wise (i.e. row j stores acceptor j's preferences).
#                                        It must be a square matrix with no zeros.
#
# O                     String     ---   Location to store the stable marriage solution matrix. If omitted, no output file is written.
#                                        If cell [i,j] is non-zero, it means that acceptor i has matched with proposer j.
#                                        Further, if cell [i,j] is non-zero, it holds the preference value that led to the match.
#
# ordered               Boolean    TRUE  If true, P and A are assumed to be ordered,
#                                        i.e. the leftmost value in a row is the most preferred partner's index.
#                                        If false, P and A are assumed to be unordered,
#                                        i.e. the leftmost value in a row in P is the preference value for the acceptor with index 1 and vice-versa (higher is better).
#
#
# Example:
# spark-submit SystemDS.jar -f StableMarriage.dml -nvargs P=Proposers.mtx A=Acceptors.mtx O=Output.mtx ordered=TRUE
#
# Proposers.mtx:
# 2.0,1.0,3.0
# 1.0,2.0,3.0
# 1.0,3.0,2.0
#
# Since ordered=TRUE, this means that proposer 1 (row 1) likes acceptor 2 the most, followed by acceptor 1 and acceptor 3.
# If ordered=FALSE, this would mean that proposer 1 (row 1) likes acceptor 3 the most (since the value at [1,3] is the row max),
# followed by acceptor 1 (2.0 preference value) and acceptor 2 (1.0 preference value).
#
# Acceptors.mtx:
# 3.0,1.0,2.0
# 2.0,1.0,3.0
# 3.0,2.0,1.0
#
# Since ordered=TRUE, this means that acceptor 1 (row 1) likes proposer 3 the most, followed by proposer 1 and proposer 2.
# If ordered=FALSE, this would mean that acceptor 1 (row 1) likes proposer 1 the most (since the value at [1,1] is the row max),
# followed by proposer 3 (2.0 preference value) and proposer 2 (1.0 preference value).
#
# Output.mtx (assuming ordered=TRUE):
# 0.0,0.0,3.0
# 0.0,3.0,0.0
# 1.0,0.0,0.0
#
# Acceptor 1 has matched with proposer 3 (since [1,3] is non-zero) at a preference level of 3.0.
# Acceptor 2 has matched with proposer 2 (since [2,2] is non-zero) at a preference level of 3.0.
# Acceptor 3 has matched with proposer 1 (since [3,1] is non-zero) at a preference level of 1.0.
#

#TODO set a finite number of maximum iterations so that the execution termiates after maximum iterations.
fileP = ifdef ($P, "");
fileA = ifdef ($A, "");
fileOutput = ifdef ($O, "");
ordered = ifdef ($ordered, TRUE);

print("\n")
print("STARTING STABLE MARRIAGE");
print("READING P AND A...");

# P : Proposers Preference Matrix
# A : Acceptors Preference Matrix

if(fileP == "" | fileA == "") {
  print("ERROR: Both P and A must be supplied.")
} 
else 
{
  P = read (fileP);
  A = read (fileA);
  if(nrow(P) != ncol(P) | nrow(A) != ncol(A)) {
    print("ERROR: Wrong Input! Both P and A must be square.")
  } 
  else if(nrow(P) != nrow(P)) {
    print("ERROR: Wrong Input! Both P and A must have the same number of rows and columns.")
  }
  else 
  {
    n = nrow(P)

    # Let S be the identity matrix
    S = diag(matrix(1.0, rows=n, cols=1))
    Result_matrix = matrix(0.0, rows=n, cols=n)

    # Pre-processing
    if(!ordered){
      # If unordered, we need to order P
      Ordered_P = matrix(0.0, rows=n, cols=n)
      Transposed_P = t(P)

      parfor(i in 1:n) 
        Ordered_P[,i] = order(target=Transposed_P, by=i, decreasing=TRUE, index.return=TRUE)
        
      P = t(Ordered_P)
    }
    else {
      # If ordered, we need to unorder A
      Unordered_A = matrix(0.0, rows=n, cols=n)

      # Since cells can be converted to unordered indices independently, we can nest parfor loops.
      parfor(i in 1:n) {
        parfor(j in 1:n, check=0) {
          Unordered_A[i, as.scalar(A[i, j])] = n - j + 1
        }
      }

      A = Unordered_A
    }

    proposer_pointers = matrix(1.0, rows=n, cols=1)

    while(sum(S) > 0) {
      Stripped_preferences = S %*% P
      Mask_matrix = matrix(0.0, rows=n, cols=n)

      parfor(i in 1:n) {
        max_proposal = as.scalar(Stripped_preferences[i, as.scalar(proposer_pointers[i])])
        if(max_proposal != 0) {
          proposer_pointers[i] = as.scalar(proposer_pointers[i]) + 1
          Mask_matrix[max_proposal, i] = 1
        }
      }

      # Hadamard Product
      Propose_round_results = Mask_matrix * A
      best_proposers_vector = rowIndexMax(Propose_round_results)
      prev_best_proposers = rowIndexMax(Result_matrix)

      parfor(i in 1:n, check=0) {
        new_best_proposer_index = as.scalar(best_proposers_vector[i])
        new_best = as.scalar(Propose_round_results[i, new_best_proposer_index])

        if(new_best > 0) {
          prev_best_proposer_index = as.scalar(prev_best_proposers[i])
          prev_best = as.scalar(Result_matrix[i, prev_best_proposer_index])

          if (new_best > prev_best) 
          {
            # Proposal is better than current fiance
            Result_matrix[i, prev_best_proposer_index] = 0
            Result_matrix[i, new_best_proposer_index] = new_best

            #Disable freshly married man to search for a new woman in the next round
            S[new_best_proposer_index, new_best_proposer_index] = 0

            # If a fiance existed, dump him/her
            if(prev_best > 0) 
              S[prev_best_proposer_index, prev_best_proposer_index] = 1
          }
        }
      }
    }

    print("Result: ")
    print(toString(Result_matrix))

    if(fileOutput != "")
      write(Result_matrix, fileOutput)
        
  }
}