import numpy as np

def compare_index(valid, valid_corr, check, check_corr):
  if np.any(valid.shape != check.shape):
    print('Output Shapes do not match')
    return False

  # Find the elements which do not agree
  no_agree = valid != check

  # SCAMP performs the best so far comparison in FP32, so if the correlations are FP32 equivalent this is okay.
  corr_fp32_no_agree = valid_corr[no_agree].astype(np.float32) != check_corr[no_agree].astype(np.float32)

  # Any indexes whose nearest neighbor correlations do not agree with fp32 precision should be treated as non-matching
  incorrect = np.count_nonzero(corr_fp32_no_agree)

  return incorrect == 0


def compare_vectors(valid, check):
  if np.any(valid.shape != check.shape):
    print('Output Shapes do not match')
    return False
    
  return np.allclose(valid, check, equal_nan=True)

def compare_vectors_sum(valid, check, thresh):
  if np.any(valid.shape != check.shape):
    print('Output Shapes do not match')
    return False
    
  x = np.allclose(valid, check, equal_nan=False)
  if not x:
    # Check if the errors are all very close to the threshold
    # Find the number of times we can fit the threshold into the difference. In the case of slight perturbation errors on results around the threshold, these should be very close to whole numbers
    diff = abs(valid - check) / thresh

    # Round to the nearest 6th decimal place (the values could be very slightly over or under the nearest whole number, we want to make sure we treat these close values as whole numbers)
    diff = np.round(diff, decimals=6)

    # Subtract the whole number part, leaving only the decimal
    diff = diff - np.floor(diff)

    # Ignore any perturbations smaller than the 6th decimal place
    diff[diff < 1e-6] = 0

    # The nonzero values in this resulting array correspond to errors that very likely do not come from perturbation errors around the threshold.
    incorrect = np.count_nonzero(diff)
    
    x = incorrect == 0
  return x

def compare_matrix(valid, check):
  if np.any(valid.shape != check.shape):
    print('Output Shapes do not match')
    return False
  
  return np.allclose(valid, check, equal_nan=True)

# This only verifies that the outputs generated by ALL_NEIGHBORS are valid 
# entries that match the true distance matrix. It does not validate that
# we found all the matches we should have found or a match we weren't supposed to find
def compare_all_neighbors(valid, check):

  incorrect = 0
  for col, row, corr in check:
    row = int(row)
    col = int(col)
    if row >= 0 and col >= 0 and row < valid.shape[0] and col < valid.shape[1]:
      if abs(valid[row, col] - corr) > matrix_match_epsilon:
        return False
    else:
      return False
  return True
