import numpy as np

def compare_index(valid, valid_corr, check, check_corr):
  if np.any(valid.shape != check.shape):
    print('Compare Index: Output Shapes valid = {} and check = {} do not match'.format(valid.shape, check.shape))
    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)

  if incorrect > 0:
    print('Compare Index: found {} incorrect indexes.'.format(incorrect))
    return False
  return True


def compare_vectors(valid, check):
  if np.any(valid.shape != check.shape):
    print('Compare Vectors: Output Shapes valid = {} and check = {} do not match'.format(valid.shape, check.shape))
    return False
  if not np.allclose(valid, check, equal_nan=True):
    print('Compare Vectors: Found inconsistencies between SCAMP output and the ground truth.')
    return False
  return True

def compare_vectors_sum(valid, check, thresh):
  if np.any(valid.shape != check.shape):
    print('Compare Sum: Output Shapes valid = {} and check = {} do not match'.format(valid.shape, check.shape))
    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)
    
    if incorrect != 0:
      print('Compare Sum: Found {} values which were far from the correct sum.'.format(incorrect))
      return False
  return True

def compare_matrix(valid_matrix, check, thresh):
  if np.any(valid_matrix.shape != check.shape):
    print('Compare Matrix: Output Shapes valid = {} and check = {} do not match'.format(valid_matrix.shape, check.shape))
    return False

  if thresh is not None:
    valid = np.copy(valid_matrix)
    valid[valid < thresh] = np.nan
  else:
    valid = valid_matrix
  
  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_matrix, valid_nearest, valid_nearest_index, check, thresh):
  rows = [int(x[1]) for x in check]
  cols = [int(x[0])for x in check]
  corrs = [x[2] for x in check]

  # All results should be above the threshold
  if np.any(np.less(corrs,thresh)):
    print('Not all results were above the threshold')
    return False

  # Duplicates should not appear.
  prev_col = -1
  row_set = set()
  nearest = np.ones(len(valid_nearest)) * -1
  nearest_index = np.ones(len(valid_nearest_index)) * -1
  for col, row, corr in zip(cols, rows, corrs):
    if col != prev_col:
      row_set = set()
      prev_col = col
      nearest[col] = corr
      nearest_index[col] = row
    if row in row_set:
      print('Duplicate row {} found in col {}'.format(row, col))
      return False
    else:
      row_set.add(row)
  
  missing_idx = nearest == -1
  valid_missing = valid_nearest[missing_idx]
  valid_missing = valid_missing[~np.isnan(valid_missing)]
  # If there are results completely missing they should have a nearest neighbor less than the threshold.
  if not np.all(valid_missing < thresh):
    print('Missing indexes were not less than the threshold')
    return False

  if (len(rows) == 0):
    return True
  
  # All outputs are present in the distance matrix somewhere
  if not np.allclose(valid_matrix[rows, cols], corrs, equal_nan=True):
    print('Not all outputs were present in the distance matrix')
    return False
    
  # Indexes which are missing due to the matches being lower than the threshold should not count against us.
  nearest[missing_idx] = valid_nearest[missing_idx] 
  nearest_index[missing_idx] = valid_nearest_index[missing_idx] 
  
  # Nearest neighbor should appear for each column. We can only guarentee the nearest neighbor will appear
  # the 2nd, 3rd ... neighbors may be different depending on how the reduction is done.
  if not compare_vectors(valid_nearest, nearest):
    printf('Nearest neighbors did not match.')
    return False 
  if not compare_index(valid_nearest_index, valid_nearest, nearest_index, nearest):
    printf('Nearest neighbor indexes did not match.')
    return False 
  return True    
