#include <Rcpp.h>


using namespace Rcpp ;
// [[Rcpp::export]]
SEXP findFrequency(NumericMatrix cooccur_matrix, int windowsize, int nas_allowed){
  int days = cooccur_matrix.ncol();
  int num_pairs = cooccur_matrix.nrow();
  NumericVector current_window(windowsize + 1);           // array that includes actual window from cooccur matrix
  int sum_contact;                                        // total number of contacts within a window
  NumericVector contacts(days);                           // number of observations of each recency in which there is a contact
  NumericVector observations(days);                       // number of observations of each recency
  

  for(int i = 0; i < num_pairs; i++) {                    // for each pair
    for(int j = 0; j < (days - windowsize); j++) {        // for each day available for window size
      for(int k = 0; k < (windowsize + 1); k++) {         // for each element in the window
        current_window[k] = cooccur_matrix(i, j + k);     // assign current window
        if(k < windowsize) {                            // if this is in the window
          if(!R_IsNA(current_window[k])) {              // if the element is not NA
            sum_contact = sum_contact + current_window[k];// calculate total number of contacts (frequency) in that window
          }
        } else {                                        // if this is the contact day after the window
          // count the number of NAs in the window
          int na_counter = 0;                           // initiate the NA counter
          for(int i = 0; i < current_window.size(); i++) { // for each element in the window
            if(R_IsNA(current_window[i])) {             // if the current element is NA
              na_counter++;                             // increment the NA counter
            }
          }
          int num_nas = na_counter;                     // assign the number of NAs to num_nas
          if(num_nas < nas_allowed + 1) {               // if there have been less than nas_allowed NAs in the window
            observations[sum_contact]++;                // increment observation counter for that frequency
            if(current_window[k] == 1) {                // if the contact day element is 1
              contacts[sum_contact]++;                  // increment the number of contacts for that frequency
            }
          }
        }
      }
      sum_contact = 0;                                    // reset number of contacts in the window
    }
  }
  return Rcpp::List::create(Rcpp::Named("freq_contacts") = contacts, Rcpp::Named("freq_observations") = observations);
}

