#include <Rmath.h>
#include <Rcpp.h>
#include "internals.h"
#include "likelihoods.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 likelihood corresponds to the probability of observing a number of
// mutations between cases and their ancestors. See src/likelihoods.cpp for
// details of the Rcpp implmentation.

// The likelihood is based on the number of mutations between a case and its
// ancestor; these are extracted from a pairwise genetic distance matrix
// (data$D) the log-likelihood is computed as: sum(mu^nmut + (1-mu)^(L-nmut))
// with:

// 'mu' is the mutation probability
// 'L' the number of sites in the alignment
// 'n_mut' the number of mutations between an ancestor and its descendent
// 'n_non_mut' the number of sites that have not mutated

// For any given case at 'nmut' mutations from its ancestor, with kappa
// generations in between, the log-likelihood is defined as:

// log(mu) * n_mut + log(1 - mu) * {(L - n_mut) + (L * (kappa-1))}


// when summing over several individuals, it becomes:

// log(mu) * sum_i(n_mut_i) + log(1-mu) * sum_i((L - n_mut_i) + (L * (kappa_i - 1)))

double cpp_ll_genetic(Rcpp::List data, Rcpp::List param, SEXP i,
		      Rcpp::RObject custom_function) {
  Rcpp::IntegerMatrix D = data["D"];
  if (D.ncol() < 1) return 0.0;

  size_t N = static_cast<size_t>(data["N"]);
  if (N < 2) return 0.0;

  Rcpp::List l;
  if (custom_function != R_NilValue) {
    l = Rcpp::as<Rcpp::List>(custom_function);
  }
  if (custom_function == R_NilValue || (l.size() > 0 && l[0] == R_NilValue)) {

    // Variables from the data & param
    Rcpp::NumericMatrix w_dens = data["log_w_dens"];
    size_t K = w_dens.nrow();
    double mu = Rcpp::as<double>(param["mu"]);
    long int L = Rcpp::as<int>(data["L"]);
    Rcpp::IntegerVector alpha = param["alpha"]; // values are on 1:N
    Rcpp::IntegerVector kappa = param["kappa"];
    Rcpp::LogicalVector has_dna = data["has_dna"];


    // Local variables used for computatoins
    size_t n_mut = 0;
    size_t n_non_mut = 0;
    double out = 0;
    bool found[1];
    size_t ances[1];
    size_t n_generations[1];
    found[0] = false;
    ances[0] = NA_INTEGER;
    n_generations[0] = NA_INTEGER;

  
    // Invalid values of mu
    if (mu < 0.0 || mu > 1.0) {
      return R_NegInf;
    }

    
    // NOTE ON MISSING SEQUENCES

    // Terms participating to the genetic likelihood correspond to pairs
    // of ancestor-descendent which have a genetic sequence. The
    // log-likelihood of other pairs is 0.0, and can therefore be
    // ommitted. Note the possible source of confusion in indices here:

    // 'has_dna' is a vector, thus indexed from 0:(N-1)
	  
    // 'cpp_get_n_mutations' is a function, and thus takes indices on 1:N


    
    // all cases are retained
    
    if (i == R_NilValue) {
      for (size_t j = 0; j < N; j++) { // 'j' on 0:(N-1)
	if (alpha[j] != NA_INTEGER) {

	  // kappa restriction

	  if (kappa[j] < 1 || kappa[j] > K) {
	    return R_NegInf;
	  }

	  // missing sequences handled here
	  
	  if (has_dna[j]) {
	 
	    lookup_sequenced_ancestor(alpha, kappa, has_dna, j + 1,
				      ances, n_generations, found);

	    if (found[0]) {

	      n_mut = cpp_get_n_mutations(data, j + 1, ances[0]); // remember the offset
	      n_non_mut = L - n_mut;

	      out += n_mut*log(n_generations[0]*mu) +
		n_non_mut*log(1 - n_generations[0]*mu);
		
	    }
	  }
	}
      }

    } else {
      // only the cases listed in 'i' are retained
      size_t length_i = static_cast<size_t>(LENGTH(i));
      Rcpp::IntegerVector vec_i(i);
      for (size_t k = 0; k < length_i; k++) {
	size_t j = vec_i[k] - 1; // offset
	if (alpha[j] != NA_INTEGER) {
	  // kappa restriction
	  if (kappa[j] < 1 || kappa[j] > K) {
	    return R_NegInf;
	  }

	  // missing sequences handled here
	  	  
	  if (has_dna[j]) {
	 
	    lookup_sequenced_ancestor(alpha, kappa, has_dna, j + 1, 
				      ances, n_generations, found);

	    if (found[0]) {

	      n_mut = cpp_get_n_mutations(data, j + 1, ances[0]); // remember the offset
	      n_non_mut = L - n_mut;
	      
	      out += n_mut*log(n_generations[0]*mu) +
		n_non_mut*log(1 - n_generations[0]*mu);
	      
	    }
	  }
	  
	}

      }
    }

    return(out);
    
  } else { // use of a customized likelihood function
    Rcpp::Function f = Rcpp::as<Rcpp::Function>(l[0]);
    int arity = l[1];
    if (arity == 3) return Rcpp::as<double>(f(data, param, i));
    return Rcpp::as<double>(f(data, param));
  }
}


double cpp_ll_genetic(Rcpp::List data, Rcpp::List param, size_t i,
              Rcpp::RObject custom_function) {
  SEXP si = PROTECT(Rcpp::wrap(i));
  double ret = cpp_ll_genetic(data, param, si, custom_function);
  UNPROTECT(1);
  return ret;
}





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

// This likelihood corresponds to the probability of observing infection dates
// of cases given the infection dates of their ancestors.

double cpp_ll_timing_infections(Rcpp::List data, Rcpp::List param, SEXP i,
				Rcpp::RObject custom_function) {
  size_t N = static_cast<size_t>(data["N"]);
  if(N < 2) return 0.0;

  Rcpp::List l;
  if (custom_function != R_NilValue) {
    l = Rcpp::as<Rcpp::List>(custom_function);
  }
  if (custom_function == R_NilValue || (l.size() > 0 && l[0] == R_NilValue)) {

    Rcpp::IntegerVector alpha = param["alpha"];
    Rcpp::IntegerVector t_inf = param["t_inf"];
    Rcpp::IntegerVector kappa = param["kappa"];
    Rcpp::NumericMatrix w_dens = data["log_w_dens"];
    size_t K = w_dens.nrow();

    double out = 0.0;

    // all cases are retained
    if (i == R_NilValue) {
      for (size_t j = 0; j < N; j++) {
	if (alpha[j] != NA_INTEGER) {
	  size_t delay = t_inf[j] - t_inf[alpha[j] - 1]; // offset
	  if (delay < 1 || delay > w_dens.ncol()) {
	    return  R_NegInf;
	  }
	  if (kappa[j] < 1 || kappa[j] > K) {
	    return  R_NegInf;
	  }

	  out += w_dens(kappa[j] - 1, delay - 1);
	}
      }
    } else {
      // only the cases listed in 'i' are retained
      size_t length_i = static_cast<size_t>(LENGTH(i));
      Rcpp::IntegerVector vec_i(i);
      for (size_t k = 0; k < length_i; k++) {
	size_t j = vec_i[k] - 1; // offset
	if (alpha[j] != NA_INTEGER) {
	  size_t delay = t_inf[j] - t_inf[alpha[j] - 1]; // offset
	  if (delay < 1 || delay > w_dens.ncol()) {
	    return  R_NegInf;
	  }
	  if (kappa[j] < 1 || kappa[j] > K) {
	    return  R_NegInf;
	  }

	  out += w_dens(kappa[j] - 1, delay - 1);
	}

      }
    }

    return out;
  } else { // use of a customized likelihood function
    Rcpp::Function f = Rcpp::as<Rcpp::Function>(l[0]);
    int arity = l[1];
    if (arity == 3) return Rcpp::as<double>(f(data, param, i));
    return Rcpp::as<double>(f(data, param));
  }
}


double cpp_ll_timing_infections(Rcpp::List data, Rcpp::List param, size_t i,
              Rcpp::RObject custom_function) {
  SEXP si = PROTECT(Rcpp::wrap(i));
  double ret = cpp_ll_timing_infections(data, param, si, custom_function);
  UNPROTECT(1);
  return ret;
}





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

// This likelihood corresponds to the probability of reporting dates of cases
// given their infection dates.

double cpp_ll_timing_sampling(Rcpp::List data, Rcpp::List param, SEXP i,
			      Rcpp::RObject custom_function) {
  size_t N = static_cast<size_t>(data["N"]);
  if(N < 2) return 0.0;

  Rcpp::List l;
  if (custom_function != R_NilValue) {
    l = Rcpp::as<Rcpp::List>(custom_function);
  }
  if (custom_function == R_NilValue || (l.size() > 0 && l[0] == R_NilValue)) {

    Rcpp::IntegerVector dates = data["dates"];
    Rcpp::IntegerVector t_inf = param["t_inf"];
    Rcpp::NumericVector f_dens = data["log_f_dens"];

    double out = 0.0;

    // all cases are retained
    if (i == R_NilValue) {
      for (size_t j = 0; j < N; j++) {
	size_t delay = dates[j] - t_inf[j];
	if (delay < 1 || delay > f_dens.size()) {
	  return  R_NegInf;
	}
	out += f_dens[delay - 1];
      }
    } else {
      // only the cases listed in 'i' are retained
      size_t length_i = static_cast<size_t>(LENGTH(i));
      Rcpp::IntegerVector vec_i(i);
      for (size_t k = 0; k < length_i; k++) {
	size_t j = vec_i[k] - 1; // offset
	size_t delay = dates[j] - t_inf[j];
	if (delay < 1 || delay > f_dens.size()) {
	  return  R_NegInf;
	}
	out += f_dens[delay - 1];
      }
    }

    return out;
  }  else { // use of a customized likelihood function
    Rcpp::Function f = Rcpp::as<Rcpp::Function>(l[0]);
    int arity = l[1];
    if (arity == 3) return Rcpp::as<double>(f(data, param, i));
    return Rcpp::as<double>(f(data, param));
  }
}


double cpp_ll_timing_sampling(Rcpp::List data, Rcpp::List param, size_t i,
              Rcpp::RObject custom_function) {
  SEXP si = PROTECT(Rcpp::wrap(i));
  double ret = cpp_ll_timing_sampling(data, param, si, custom_function);
  UNPROTECT(1);
  return ret;
}





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

// This likelihood corresponds to the probability of a given number of
// unreported cases on an ancestry.

// The likelihood is given by a geometric distribution with probability 'pi'
// to report a case

// - 'kappa' is the number of generation between two successive cases
// - 'kappa-1' is the number of unreported cases

double cpp_ll_reporting(Rcpp::List data, Rcpp::List param, SEXP i,
			Rcpp::RObject custom_function) {
  Rcpp::NumericMatrix w_dens = data["log_w_dens"];
  size_t K = w_dens.nrow();

  size_t N = static_cast<size_t>(data["N"]);
  if(N < 2) return 0.0;

  double pi = static_cast<double>(param["pi"]);
  Rcpp::IntegerVector kappa = param["kappa"];

  // p(pi < 0) = p(pi > 1) = 0
  if (pi < 0.0 || pi > 1.0) {
    return R_NegInf;
  }

  Rcpp::List l;
  if (custom_function != R_NilValue) {
    l = Rcpp::as<Rcpp::List>(custom_function);
  }
  if (custom_function == R_NilValue || (l.size() > 0 && l[0] == R_NilValue)) {

    double out = 0.0;

    // all cases are retained
    if (i == R_NilValue) {
      for (size_t j = 0; j < N; j++) {
	if (kappa[j] != NA_INTEGER) {
	  if (kappa[j] < 1 || kappa[j] > K) {
	    return  R_NegInf;
	  }
	  out += R::dgeom(kappa[j] - 1.0, pi, 1); // first arg must be cast to double
	}
      }
    } else {
      // only the cases listed in 'i' are retained
      size_t length_i = static_cast<size_t>(LENGTH(i));
      Rcpp::IntegerVector vec_i(i);
      for (size_t k = 0; k < length_i; k++) {
	size_t j = vec_i[k] - 1; // offset
	if (kappa[j] != NA_INTEGER) {
	  if (kappa[j] < 1 || kappa[j] > K) {
	    return  R_NegInf;
	  }
	  out += R::dgeom(kappa[j] - 1.0, pi, 1); // first arg must be cast to double
	}
      }
    }

    return out;
  } else { // use of a customized likelihood function
    Rcpp::Function f = Rcpp::as<Rcpp::Function>(l[0]);
    int arity = l[1];
    if (arity == 3) return Rcpp::as<double>(f(data, param, i));
    return Rcpp::as<double>(f(data, param));
  }
}


double cpp_ll_reporting(Rcpp::List data, Rcpp::List param, size_t i,
              Rcpp::RObject custom_function) {
  SEXP si = PROTECT(Rcpp::wrap(i));
  double ret = cpp_ll_reporting(data, param, si, custom_function);
  UNPROTECT(1);
  return ret;
}





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

// This likelihood corresponds to the probability of observing a a reported
// contact between cases and their ancestors. See
// src/likelihoods.cpp for details of the Rcpp implmentation.

// The likelihood is based on the contact status between a case and its
// ancestor; this is extracted from a pairwise contact matrix (data$C), the
// log-likelihood is computed as:
// true_pos*eps + false_pos*eps*xi +
// false_neg*(1- eps) + true_neg*(1 - eps*xi)
//
// with:
// 'eps' is the contact reporting coverage
// 'lambda' is the non-infectious contact rate
// 'true_pos' is the number of contacts between transmission pairs
// 'false_pos' is the number of contact between non-transmission pairs
// 'false_neg' is the number of transmission pairs without contact
// 'true_neg' is the number of non-transmission pairs without contact

double cpp_ll_contact(Rcpp::List data, Rcpp::List param, SEXP i,
		       Rcpp::RObject custom_function) {
  Rcpp::NumericMatrix contacts = data["contacts"];
  if (contacts.ncol() < 1) return 0.0;

  size_t C_combn = static_cast<size_t>(data["C_combn"]);
  size_t C_nrow = static_cast<size_t>(data["C_nrow"]);

  size_t N = static_cast<size_t>(data["N"]);
  if (N < 2) return 0.0;

  Rcpp::List l;
  if (custom_function != R_NilValue) {
    l = Rcpp::as<Rcpp::List>(custom_function);
  }
  if (custom_function == R_NilValue || (l.size() > 0 && l[0] == R_NilValue)) {

    double out;
    double eps = Rcpp::as<double>(param["eps"]);
    double lambda = Rcpp::as<double>(param["lambda"]);
    Rcpp::IntegerVector alpha = param["alpha"];
    Rcpp::IntegerVector kappa = param["kappa"];

    size_t true_pos = 0;
    size_t false_pos = 0;
    size_t false_neg = 0;
    size_t true_neg = 0;
    size_t imports = 0;
    size_t unobsv_case = 0;

    // p(eps < 0 || lambda < 0) = 0
    if (eps < 0.0 || lambda < 0.0) {
      return R_NegInf;
    }

    // all cases are retained (currently no support for i subsetting)
    for (size_t j = 0; j < N; j++) {
      if (alpha[j] == NA_INTEGER) {
	imports += 1;
      } else if (kappa[j] > 1) {
	unobsv_case += 1;
      } else {
	true_pos += contacts(j, alpha[j] - 1); // offset
      }
    }

    false_pos = C_nrow - true_pos;
    false_neg = N - imports - unobsv_case - true_pos;
    true_neg = C_combn - true_pos - false_pos - false_neg;

    // deal with special case when lambda == 0 and eps == 1, to avoid log(0)
    if(lambda == 0.0) {
      if(false_pos > 0) {
	out = R_NegInf;
      } else {
	out = log(eps) * (double) true_pos +
	  log(1 - eps) * (double) false_neg +
	  log(1 - eps*lambda) * (double) true_neg;
      }
    } else if(eps == 1.0) {
      if(false_neg > 0) {
	out = R_NegInf;
      } else {
	out = log(eps) * (double) true_pos +
	  log(eps*lambda) * (double) false_pos +
	  log(1 - eps*lambda) * (double) true_neg;
      }
    } else {
      out = log(eps) * (double) true_pos +
	log(eps*lambda) * (double) false_pos +
	log(1 - eps) * (double) false_neg +
	log(1 - eps*lambda) * (double) true_neg;
    }

    return out;
    
  } else { //use of a customized likelihood function
    Rcpp::Function f = Rcpp::as<Rcpp::Function>(l[0]);
    int arity = l[1];
    if (arity == 3) return Rcpp::as<double>(f(data, param, i));
    return Rcpp::as<double>(f(data, param));
  }
}


double cpp_ll_contact(Rcpp::List data, Rcpp::List param, size_t i,
              Rcpp::RObject custom_function) {
  SEXP si = PROTECT(Rcpp::wrap(i));
  double ret = cpp_ll_contact(data, param, si, custom_function);
  UNPROTECT(1);
  return ret;
}





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

// This likelihood corresponds to the sums of the separate timing likelihoods,
// which include:

// - p(infection dates): see function cpp_ll_timing_infections
// - p(collection dates): see function cpp_ll_timing_sampling

double cpp_ll_timing(Rcpp::List data, Rcpp::List param, SEXP i,
		     Rcpp::RObject custom_functions) {

  if (custom_functions == R_NilValue) {
    return cpp_ll_timing_infections(data, param, i) +
      cpp_ll_timing_sampling(data, param, i);
  } else { // use of a customized likelihood functions
    Rcpp::List list_functions = Rcpp::as<Rcpp::List>(custom_functions);
    return cpp_ll_timing_infections(data, param, i, list_functions["timing_infections"]) +
      cpp_ll_timing_sampling(data, param, i, list_functions["timing_sampling"]);

  }
}


double cpp_ll_timing(Rcpp::List data, Rcpp::List param, size_t i,
              Rcpp::RObject custom_function) {
  SEXP si = PROTECT(Rcpp::wrap(i));
  double ret = cpp_ll_timing(data, param, si, custom_function);
  UNPROTECT(1);
  return ret;
}




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

// This likelihood corresponds to the sums of the separate likelihoods, which
// include:

// - p(infection dates): see function cpp_ll_timing_infections
// - p(collection dates): see function cpp_ll_timing_sampling
// - p(genetic diversity): see function cpp_ll_genetic
// - p(missing cases): see function cpp_ll_reporting
// - p(contact): see function cpp_ll_contact 

double cpp_ll_all(Rcpp::List data, Rcpp::List param, SEXP i,
		  Rcpp::RObject custom_functions) {

  if (custom_functions == R_NilValue) {

    return cpp_ll_timing_infections(data, param, i) +
      cpp_ll_timing_sampling(data, param, i) +
      cpp_ll_genetic(data, param, i) +
      cpp_ll_reporting(data, param, i) +
      cpp_ll_contact(data, param, i);

  }  else { // use of a customized likelihood functions
    Rcpp::List list_functions = Rcpp::as<Rcpp::List>(custom_functions);

    return cpp_ll_timing_infections(data, param, i, list_functions["timing_infections"]) +
      cpp_ll_timing_sampling(data, param, i, list_functions["timing_sampling"]) +
      cpp_ll_genetic(data, param, i, list_functions["genetic"]) +
      cpp_ll_reporting(data, param, i, list_functions["reporting"]) +
      cpp_ll_contact(data, param, i, list_functions["contact"]);

  }
}


double cpp_ll_all(Rcpp::List data, Rcpp::List param, size_t i,
              Rcpp::RObject custom_function) {
  SEXP si = PROTECT(Rcpp::wrap(i));
  double ret = cpp_ll_all(data, param, si, custom_function);
  UNPROTECT(1);
  return ret;
}