#include "internals.h"



// IMPORTANT: ON INDEXING VECTORS AND ANCESTRIES

// Most of the functions implemented here are susceptible to be called from R
// via Rcpp, and are therefore treated as interfaces. This causes a number of
// headaches when using indices of cases defined in R (1:N) to refer to elements
// in Rcpp / Cpp vectors (0:N-1). By convention, we store all data on the
// original scale (1:N), and modify indices whenever accessing elements of
// vectors. In other words, in an expression like 'alpha[j]', 'j' should always
// be on the internal scale (0:N-1).

// In all these functions, 'SEXP i' is an optional vector of case indices, on
// the 1:N scale.





// ---------------------------

//   This function returns a vector of indices of cases which could be infector
//   of 'i' (i.e., their infection dates preceed that of 'i'). Only tricky bit
//   here is keep in mind that 't_inf' is indexed from 0 to N-1, while 'i' and
//   'alpha' (ancestors) are values from 1 to N.

//   Original R code:

// are.possible.alpha <- function(t_inf, i) {
//     if (length(i)>1) {
//         stop("i has a length > 1")
//     }
//     if (any(t_inf[i]==min(t_inf))) {
//         return(NA)
//     }
//     return(which(t_inf < t_inf[i[1]]))
// }

// [[Rcpp::export()]]
std::vector<int> cpp_are_possible_ancestors(Rcpp::IntegerVector t_inf, size_t i) {
  size_t n = t_inf.size();
  std::vector<int> out;
  out.reserve(n);
  for (size_t j = 0; j < n; j++) {
    if (t_inf[j] < t_inf[i-1]) { // offset
      out.push_back(j+1); // +1 needed for range 1 ... N
    }
  }
  return out;
}





// ---------------------------

//  This function samples a single value from a vector of integers.

// [[Rcpp::export()]]
size_t cpp_sample1(std::vector<int> x) {
  if (x.size() < 1) {
    Rcpp::Rcerr << "Trying to sample from empty vector" << std::endl;
    Rcpp::stop("Trying to sample from empty vector");
  }

  return x[unif_rand() * x.size()];
}




// ---------------------------

//    This function choose a possible infector for case 'i'; 'i' is on the scale
//    1:N

//    Original R version:

// .choose.possible.alpha <- function(t_inf, i) {
//     return(sample(are.possible.alpha(t_inf=t_inf, i=i), 1))
// }

// [[Rcpp::export()]]
size_t cpp_pick_possible_ancestor(Rcpp::IntegerVector t_inf, size_t i) {
  return cpp_sample1(cpp_are_possible_ancestors(t_inf, i));
}






// ---------------------------

// This function returns the descendents of a given case 'i' in the current
// ancestries; 'i' is on the scale 1:N. The output is also on the scale 1:N.

// Original R version:

// find.descendents <- function(param, i) {
//   ## find descendents
//     which(param.current$alpha==i)
//  }

// [[Rcpp::export()]]
Rcpp::IntegerVector cpp_find_descendents(Rcpp::IntegerVector alpha, size_t i) {
  size_t counter = 0, n = 0;

  // determine size of output vector and create it
  for (size_t j = 0; j < alpha.size(); j++) {
    if (alpha[j] == i) n++;
  }
  
  Rcpp::IntegerVector out(n);

  // fill in output vector
  for (size_t j = 0; j < alpha.size(); j++) {
    if (alpha[j] == i) {
      out[counter++] = j + 1; // offset
    }
  }
  return out;
}







// ---------------------------

// This function returns a vector of indices of cases which are 'local' to a
// case 'i'. Locality is defined as the following set of cases:

// - 'i'
// - the descendents of 'i'
// - 'alpha[i-1]'
// - the descendents of 'alpha[i]' (excluding 'i')

// where 'alpha' is a IntegerVector storing ancestries. Note that 'i' and
// 'alpha' are on the scale 1:N. 

// [[Rcpp::export()]]
Rcpp::IntegerVector cpp_find_local_cases(Rcpp::IntegerVector alpha, size_t i) {
  // determine descendents of 'i':
  Rcpp::IntegerVector desc_i = cpp_find_descendents(alpha, i);
  size_t n = desc_i.size() + 1; // +1 is to count 'i' itself
  
  // determine descendents of 'alpha[i]':
  Rcpp::IntegerVector desc_alpha_i = cpp_find_descendents(alpha,
							  (size_t) alpha[i-1]);
  if (alpha[i-1] != NA_INTEGER) {
    n += desc_alpha_i.size();
  }

  // create output
  Rcpp::IntegerVector out(n);
  size_t counter = 0;

  // 'i'
  out[counter++] = i;

  // 'descendents of 'i'
  for (size_t j = 0; j < desc_i.size(); j++) {
    out[counter++] = desc_i[j];
  }
  
  if (alpha[i-1] != NA_INTEGER) {
    // alpha[i-1] ...
    out[counter++] = alpha[i-1];
    
    // ... and its descendents
    for (size_t j = 0; j < desc_alpha_i.size(); j++) {
      if ( desc_alpha_i[j] != i) {
	out[counter++] = desc_alpha_i[j];
      }
    }
  }

  return out;
}








// ---------------------------

// This function swaps cases in a transmission tree. The focus case is 'i', and
// is swapped with its ancestor 'x=alpha[i-1]'. In other words the change is
// from: x -> i to i -> x
// Involved changes are:

// - descendents of 'i' become descendents of 'x'
// - descendents of 'x' become descendents of 'i'
// - the infector if 'i' becomes the infector of 'x' (i.e. alpha[x-1])
// - the infector if 'x' becomes 'i'
// - infection time of 'i' becomes that of 'x'
// - infection time of 'x' becomes that of 'i'

// Note on indexing: 'i', 'x', and values of alpha are on the scale 1:N. The
// function's output is a list with new alpha and t_inf.

// Note on forbidden swaps: two types of swaps are excluded:
// - 'i' is imported, so that 'alpha[i-1]' is NA_INTEGER
// - 'x' is imported, so that 'alpha[x-1]' is NA_INTEGER

// [[Rcpp::export()]]
Rcpp::List cpp_swap_cases(Rcpp::List param, size_t i) {
  Rcpp::IntegerVector alpha_in = param["alpha"];
  Rcpp::IntegerVector t_inf_in = param["t_inf"];
  Rcpp::IntegerVector kappa_in = param["kappa"];
  Rcpp::IntegerVector alpha_out = clone(alpha_in);
  Rcpp::IntegerVector t_inf_out = clone(t_inf_in);
  Rcpp::IntegerVector kappa_out = clone(kappa_in);
  Rcpp::List out;
  out["alpha"] = alpha_out;
  out["t_inf"] = t_inf_out;
  out["kappa"] = kappa_out;
      
  size_t N = alpha_in.size();
  
  // escape if the case is imported, i.e. alpha[i-1] is NA
  
  if (alpha_in[i-1] == NA_INTEGER) {
    return out;
  }


  // escape if ancestor of the case is imported, i.e. alpha[x-1] is NA
  
  size_t x = (size_t) alpha_in[i-1];
  //if (alpha_in[x-1] == NA_INTEGER) {
  // return out;
  //}
  
 
  // replace ancestries:
  // - descendents of 'i' become descendents of 'x'
  // - descendents of 'x' become descendents of 'i'

  for (size_t j = 0; j < N; j++) {
    if (alpha_in[j] == i) {
      alpha_out[j] = x;
    } else if (alpha_in[j] == x) {
      alpha_out[j] = i;
    }
  }


  // the ancestor of 'i' becomes an ancestor of 'x'

  alpha_out[i-1] = alpha_in[x-1];

  
  // 'i' is now the ancestor of 'x'
  alpha_out[x-1] = i;
  

  // swap infections times of 'i' and 'x'
  t_inf_out[i-1] =   t_inf_in[x-1];
  t_inf_out[x-1] =   t_inf_in[i-1];

  kappa_out[i-1] =   kappa_in[x-1];
  kappa_out[x-1] =   kappa_in[i-1];


  return out;
}






// ---------------------------

// This function returns the number of mutations between two cases from a 'data'
// object. It uses the indexing of cases in the DNA matrix to ensure
// correspondance between cases and their sequences (not all cases may have a
// sequence). 

// i and j are indices of cases on the scale 1:N; note that the vectors and
// matrices are indexed on 0:(N-1).

// [[Rcpp::export()]]
size_t cpp_get_n_mutations(Rcpp::List data, size_t i, size_t j) {
  Rcpp::LogicalVector has_dna = data["has_dna"];
  Rcpp::IntegerVector id_in_dna = data["id_in_dna"];
  Rcpp::IntegerMatrix D = data["D"];

  
  // Ideally we should return NA_integer here, but then the type of the function
  // should be a Rcpp::IntegerVector, which would complicate things. The second
  // best thing we can do here really is to issue an error when trying to get
  // number of mutations between cases with missing sequences.
  
  if (!(has_dna[i-1] && has_dna[j-1])) {
    Rcpp::stop("Trying to get genetic distances between missing sequences.");
  } 

  size_t out = D(id_in_dna[i-1] - 1, id_in_dna[j-1] - 1);

  return out;
  
}






// ---------------------------

// This function looks up a transmission chain to find the most recent ancestor
// with a sequence, for a given case 'i'. It stops at two conditions: i) it
// finds a sequenced ancestor, or ii) the current ancestor is 'NA'. It returns a
// List with three values: i) the index of the most recent ancestor (on the
// scale 1:N), ii) the total number of generations between this case and 'i',
// and iii) a logical 'found_sequenced_ancestor'. If the latter is FALSE, then
// previous values are 'NA_INTEGER'.

// This is the exported interface. It calls upon a non-exported function
// (lookup_sequenced_ancestor) which does not make memory allocation for the
// output, but instead modifies one of its arguments. This trade-off pays as it
// allows for unit testing via the interface, but remains quite fast as the
// non-exported function can be used internally. "i" is indexed on 1:N.

// [[Rcpp::export()]]
Rcpp::List cpp_lookup_sequenced_ancestor(Rcpp::List data, Rcpp::List param, size_t i) {
 
  Rcpp::IntegerVector alpha = param["alpha"];
  Rcpp::IntegerVector kappa = param["kappa"];
  Rcpp::LogicalVector has_dna = data["has_dna"];
  
  Rcpp::List out;
  Rcpp::IntegerVector out_ances(1);
  Rcpp::IntegerVector out_n_generations(1);
  Rcpp::LogicalVector out_found_sequenced_ancestor(1);
  out["alpha"] = out_ances;
  out["n_generations"] = out_n_generations;
  out["found_sequenced_ancestor"] = out_found_sequenced_ancestor;
  
  size_t ances[1];
  size_t n_generations[1];
  bool found_sequenced_ancestor[1];

  ances[0] = NA_INTEGER;
  n_generations[0] = NA_INTEGER;
  found_sequenced_ancestor[0] = false;

  // This function modifies its last argument
  lookup_sequenced_ancestor(alpha, kappa, has_dna, i, // inputs
			    ances, n_generations, 
			    found_sequenced_ancestor); // outputs


  out_ances[0] = static_cast<int>(ances[0]);
  out_n_generations[0] =  static_cast<int>(n_generations[0]); 
  out_found_sequenced_ancestor[0] = found_sequenced_ancestor[0];

  return out;
}






// ---------------------------

// This function is the internal version of cpp_lookup_sequenced_ancestor. It is
// not meant to be called by users, only by internal procedures, as it modifies
// the content of its last argument rather than creating a new object, which is
// obviously dangerous. Only use it carefully if you handled the creating of its
// last argument 'out'. 'out_' are technically outputs with three components:
// "ances" (IntegerVector of size 1), "n_generations" (same), and
// "found_sequenced_ancestor" (LogicalVector of length 1). "i" is indexed on
// 1:N.

void lookup_sequenced_ancestor(Rcpp::IntegerVector alpha, Rcpp::IntegerVector kappa, 
			       Rcpp::LogicalVector has_dna, size_t i, 
			       size_t *out_alpha, 
			       size_t *out_n_generations, 
			       bool *out_found_sequenced_ancestor
			       ) {

  if (!has_dna[i - 1] || alpha[i - 1] == NA_INTEGER) {
    return;
  }
  
 
  size_t current_case = i; // this one is indexed on 1:N
  size_t n_generations = kappa[current_case - 1];
  bool ances_has_dna = has_dna[alpha[current_case - 1] - 1]; // offset for indexing vectors
    

  // look recursively for ancestor with sequence if needed
	    
  while (!ances_has_dna && (alpha[current_case - 1] != NA_INTEGER)) {
    current_case = alpha[current_case - 1]; // 1 step back up the transmission chain
    ances_has_dna = (alpha[current_case - 1] != NA_INTEGER) && // need to test for NA *first*
      has_dna[alpha[current_case - 1] - 1]; // offset for indexing vectors
    n_generations += kappa[current_case - 1];
  }


  // change outputs as needed

  
  if (ances_has_dna) {
      out_alpha[0] =  alpha[current_case - 1];
      out_n_generations[0] = n_generations;
      out_found_sequenced_ancestor[0] = true;
  } else {
    out_alpha[0] = NA_INTEGER;
    out_n_generations[0] = NA_INTEGER;
    out_found_sequenced_ancestor[0] = false;
  }
  
}