#include <Rcpp.h>
#include <Rmath.h>
#include "internals.h"
#include "likelihoods.h"
#include "priors.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.






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

// Movement of the mutation rate 'mu' is done using a dumb normal proposal. This
// is satisfying for now - we only reject a few non-sensical values outside the
// range [0;1]. The SD of the proposal (implicitely contained in rand$mu.rnorm1,
// but really provided through 'config', seems fine as the range of real values
// will never change much. Probably not much point in using auto-tuning here.

// [[Rcpp::export(rng = true)]]
Rcpp::List cpp_move_mu(Rcpp::List param, Rcpp::List data, Rcpp::List config,
		       Rcpp::RObject custom_ll = R_NilValue,
		       Rcpp::RObject custom_prior = R_NilValue) {

  // deep copy here for now, ultimately should be an arg.
  Rcpp::List new_param = clone(param);
  Rcpp::NumericVector mu = param["mu"];
  Rcpp::NumericVector new_mu = new_param["mu"];

  double sd_mu = static_cast<double>(config["sd_mu"]);

  double old_logpost = 0.0, new_logpost = 0.0, p_accept = 0.0;

  // proposal (normal distribution with SD: config$sd_mu)

  new_mu[0] += R::rnorm(0.0, sd_mu); // new proposed value


  // automatic rejection of negative mu

  if (new_mu[0] < 0.0) {
    return param;
  }


  // compute likelihoods
  old_logpost = cpp_ll_genetic(data, param, R_NilValue, custom_ll);
  new_logpost = cpp_ll_genetic(data, new_param, R_NilValue, custom_ll);


  // compute priors

  old_logpost += cpp_prior_mu(param, config, custom_prior);
  new_logpost += cpp_prior_mu(new_param, config, custom_prior);


  // acceptance term

  p_accept = exp(new_logpost - old_logpost);


  // acceptance: the new value is already in mu, so we only act if the move is
  // rejected, in which case we restore the previous ('old') value

  if (p_accept < unif_rand()) { // reject new values
    return param;
  }

  return new_param;
}






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

// movement of the Reporting probability 'pi' is done using a dumb normal
// proposal. This is satisfying for now - we only reject a few non-sensical
// values outside the range [0;1]. The SD of the proposal (implicitely contained
// in rand$pi.rnorm1, but really provided through 'config', seems fine as the
// range of real values will never change much. Probably not much point in using
// auto-tuning here.

// [[Rcpp::export(rng = true)]]
Rcpp::List cpp_move_pi(Rcpp::List param, Rcpp::List data, Rcpp::List config,
		       Rcpp::RObject custom_ll = R_NilValue,
		       Rcpp::RObject custom_prior = R_NilValue) {

  // deep copy here for now, ultimately should be an arg.

  Rcpp::List new_param = clone(param);
  Rcpp::NumericVector pi = param["pi"]; // these are just pointers
  Rcpp::NumericVector new_pi = new_param["pi"]; // these are just pointers

  double sd_pi = static_cast<double>(config["sd_pi"]);

  double old_logpost = 0.0, new_logpost = 0.0, p_accept = 0.0;


  // proposal (normal distribution with SD: config$sd_pi)

  new_pi[0] += R::rnorm(0.0, sd_pi); // new proposed value


  // automatic rejection of pi outside [0;1]

  if (new_pi[0] < 0.0 || new_pi[0] > 1.0) {
    return param;
  }


  // compute likelihoods
  old_logpost = cpp_ll_reporting(data, param, R_NilValue, custom_ll);
  new_logpost = cpp_ll_reporting(data, new_param, R_NilValue, custom_ll);


  // compute priors

  old_logpost += cpp_prior_pi(param, config, custom_prior);
  new_logpost += cpp_prior_pi(new_param, config, custom_prior);


  // acceptance term

  p_accept = exp(new_logpost - old_logpost);


  // acceptance: the new value is already in pi, so we only act if the move is
  // rejected, in which case we restore the previous ('old') value

  if (p_accept < unif_rand()) { // reject new values
    return param;
  }

  return new_param;
}






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

// movement of the contact reporting coverage 'eps' is done using a dumb normal
// proposal. This is satisfying for now - we only reject a few non-sensical
// values outside the range [0;1]. The SD of the proposal is provided through
// 'config'; this seems fine as the range of real values will never change
// much. Probably not much point in using auto-tuning here.

// [[Rcpp::export(rng = true)]]
Rcpp::List cpp_move_eps(Rcpp::List param, Rcpp::List data, Rcpp::List config,
		       Rcpp::RObject custom_ll = R_NilValue,
		       Rcpp::RObject custom_prior = R_NilValue) {

  // deep copy here for now, ultimately should be an arg.

  Rcpp::List new_param = clone(param);
  Rcpp::NumericVector eps = param["eps"]; // these are just pointers
  Rcpp::NumericVector new_eps = new_param["eps"]; // these are just pointers

  double sd_eps = static_cast<double>(config["sd_eps"]);

  double old_logpost = 0.0, new_logpost = 0.0, p_accept = 0.0;


  // proposal (normal distribution with SD: config$sd_eps)

  new_eps[0] += R::rnorm(0.0, sd_eps); // new proposed value


  // automatic rejection of eps outside [0;1]

  if (new_eps[0] < 0.0 || new_eps[0] > 1.0) {
    return param;
  }


  // compute likelihoods
  old_logpost = cpp_ll_contact(data, param, R_NilValue, custom_ll);
  new_logpost = cpp_ll_contact(data, new_param, R_NilValue, custom_ll);


  // compute priors

  old_logpost += cpp_prior_eps(param, config, custom_prior);
  new_logpost += cpp_prior_eps(new_param, config, custom_prior);


  // acceptance term

  p_accept = exp(new_logpost - old_logpost);


  // acceptance: the new value is already in eps, so we only act if the move is
  // rejected, in which case we restore the previous ('old') value

  if (p_accept < unif_rand()) { // reject new values
    return param;
  }

  return new_param;
}






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

// movement of the non-infectious contact rate 'eps' is done using a dumb
// normal proposal. This is satisfying for now - we only reject a few
// non-sensical values outside the range [0;1]. The SD of the proposal is
// provided through 'config'; this seems fine as the range of real values will
// never change much. Probably not much point in using auto-tuning here.

// [[Rcpp::export(rng = true)]]
Rcpp::List cpp_move_lambda(Rcpp::List param, Rcpp::List data, Rcpp::List config,
		       Rcpp::RObject custom_ll = R_NilValue,
		       Rcpp::RObject custom_prior = R_NilValue) {

  // deep copy here for now, ultimately should be an arg.

  Rcpp::List new_param = clone(param);
  Rcpp::NumericVector lambda = param["lambda"]; // these are just pointers
  Rcpp::NumericVector new_lambda = new_param["lambda"]; // these are just pointers

  double sd_lambda = static_cast<double>(config["sd_lambda"]);

  double old_logpost = 0.0, new_logpost = 0.0, p_accept = 0.0;


  // proposal (normal distribution with SD: config$sd_lambda)

  new_lambda[0] += R::rnorm(0.0, sd_lambda); // new proposed value


  // automatic rejection of lambda outside [0;1]

  if (new_lambda[0] < 0.0 || new_lambda[0] > 1.0) {
    return param;
  }


  // compute likelihoods
  old_logpost = cpp_ll_contact(data, param, R_NilValue, custom_ll);
  new_logpost = cpp_ll_contact(data, new_param, R_NilValue, custom_ll);


  // compute priors

  old_logpost += cpp_prior_lambda(param, config, custom_prior);
  new_logpost += cpp_prior_lambda(new_param, config, custom_prior);


  // acceptance term

  p_accept = exp(new_logpost - old_logpost);


  // acceptance: the new value is already in lambda, so we only act if the move is
  // rejected, in which case we restore the previous ('old') value

  if (p_accept < unif_rand()) { // reject new values
    return param;
  }

  return new_param;
}






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

// Movement of infection dates are +/- 1 from current states. These movements
// are currently vectorised, i.e. a bunch of dates are proposed all together;
// this may not be sustainable for larger datasets. The non-vectorised option
// will be slower and speed-up with C/C++ will be more substantial then.

// This version differs from the initial R implementation in several points:

// 1. all cases are moved
// 2. cases are moved one by one
// 3. movement for each case is +/- 1 time unit

// Notes

// - when computing the timing log-likelihood, the descendents of each
// case are also affected.

// - we generate a new vector 'new_t_inf', which will replace the
// previous pointer defining param["t_inf"].

// [[Rcpp::export(rng = true)]]
Rcpp::List cpp_move_t_inf(Rcpp::List param, Rcpp::List data,
			  Rcpp::RObject list_custom_ll = R_NilValue) {

  // deep copy here for now, ultimately should be an arg.

  Rcpp::List new_param = clone(param);
  Rcpp::IntegerVector t_inf = param["t_inf"];
  Rcpp::IntegerVector new_t_inf = new_param["t_inf"];
  Rcpp::IntegerVector alpha = param["alpha"];
  Rcpp::IntegerVector local_cases;

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

  double old_loc_loglike = 0.0, new_loc_loglike = 0.0, p_loc_accept = 0.0;


  for (size_t i = 0; i < N; i++) {

    local_cases = cpp_find_descendents(param["alpha"], i+1);

    // loglike with current value
    old_loc_loglike = cpp_ll_timing(data, param, i+1, list_custom_ll); // term for case 'i' with offset

    // term descendents of 'i'
    if (local_cases.size() > 0) {
      old_loc_loglike += cpp_ll_timing(data, param, local_cases, list_custom_ll);
    }

    // proposal (+/- 1)
    new_t_inf[i] += unif_rand() > 0.5 ? 1 : -1; // new proposed value

    // loglike with new value
    new_loc_loglike = cpp_ll_timing(data, new_param, i+1, list_custom_ll); // term for case 'i' with offset

    // term descendents of 'i'
    if (local_cases.size() > 0) {
      new_loc_loglike += cpp_ll_timing(data, new_param, local_cases, list_custom_ll);
    }


    // acceptance term
    p_loc_accept = exp(new_loc_loglike - old_loc_loglike);


    // acceptance: the new value is already in t_inf, so we only act if the move
    // is rejected, in which case we restore the previous ('old') value

    if (p_loc_accept < unif_rand()) { // reject new values
      new_t_inf[i] = t_inf[i];
    }
  }

  return new_param;

}






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

// Movement of ancestries ('alpha') is not vectorised, movements are made one
// case at a time. This procedure is simply about picking an infector at random
// amongst cases preceeding the case considered. The original version in
// 'outbreaker' used to move simultaneously 'alpha', 'kappa' and 't_inf', but
// current implementation is simpler and seems to mix at least as well. Proper
// movement of 'alpha' needs this procedure as well as a swapping procedure
// (swaps are not possible through move.alpha only); in all instances, 'alpha'
// is on the scale 1:N.

// [[Rcpp::export(rng = true)]]
Rcpp::List cpp_move_alpha(Rcpp::List param, Rcpp::List data,
			  Rcpp::RObject list_custom_ll = R_NilValue) {
  Rcpp::List new_param = clone(param);
  Rcpp::IntegerVector alpha = param["alpha"]; // pointer to param$alpha
  Rcpp::IntegerVector t_inf = param["t_inf"]; // pointer to param$t_inf
  Rcpp::IntegerVector new_alpha = new_param["alpha"];

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

  double old_loglike = 0.0, new_loglike = 0.0, p_accept = 0.0;


  for (size_t i = 0; i < N; i++) {

    // only non-NA ancestries are moved, if there is at least 1 choice
    if (alpha[i] != NA_INTEGER && sum(t_inf < t_inf[i]) > 0) {

      // loglike with current value
      // old_loglike = cpp_ll_all(data, param, R_NilValue);
      old_loglike = cpp_ll_all(data, param, i+1, list_custom_ll); // offset

      // proposal (+/- 1)
      new_alpha[i] = cpp_pick_possible_ancestor(t_inf, i+1); // new proposed value (on scale 1 ... N)

      // loglike with current value
      new_loglike = cpp_ll_all(data, new_param, i+1, list_custom_ll);

      // acceptance term
      p_accept = exp(new_loglike - old_loglike);

      // which case we restore the previous ('old') value
      if (p_accept < unif_rand()) { // reject new values
	new_alpha[i] = alpha[i];
      } else {
	alpha[i] = new_alpha[i];
      }
    }
  }

  return new_param;
}






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

// The basic movement of ancestries (picking an ancestor at random amongst in
// previous cases) makes swaps of ancestries (A->B) to (B->A) very
// difficult. This function addresses the issue. It is computer-intensive, but
// likely a determining factor for faster mixing. Unlike previous versions in
// the original 'outbreaker' package, all cases are 'moved' here. A swap is
// defined as:

// x -> a -> b  becomes a -> x -> b

// Obviously cases are moved one at a time. We need to use local likelihood
// changes for this move to scale well with outbreak size. The complicated bit
// is that the move impacts all descendents from 'a' as well as 'x'.

// [[Rcpp::export(rng = true)]]
Rcpp::List cpp_move_swap_cases(Rcpp::List param, Rcpp::List data,
			       Rcpp::RObject list_custom_ll = R_NilValue) {

  Rcpp::List new_param = clone(param);
  Rcpp::IntegerVector alpha = param["alpha"]; // pointer to param$alpha
  Rcpp::IntegerVector t_inf = param["t_inf"]; // pointer to param$t_inf
  Rcpp::List swapinfo; // contains alpha and t_inf
  Rcpp::IntegerVector local_cases;

  size_t N = static_cast<size_t>(data["N"]);
  int N_int = data["N"];
  
  double old_loglike = 0.0, new_loglike = 0.0, p_accept = 0.0;

  // Shuffle indices to make equal cases equally likely
  Rcpp::IntegerVector idx = Rcpp::seq(0, N-1);
  idx = Rcpp::sample(idx, N_int, false);

  for (size_t j = 0; j < N; ++j) {
    size_t i = (size_t)idx[j];

    // only non-NA ancestries are moved, if there is at least 1 choice
    if (alpha[i] != NA_INTEGER && sum(t_inf < t_inf[i]) > 0) {

      // The local likelihood is defined as the likelihood computed for the
      // cases affected by the swap; these include:

      // - 'i'
      // - the descendents of 'i'
      // - 'alpha[i]'
      // - the descendents of 'alpha[i]' (other than 'i')

      local_cases = cpp_find_local_cases(param["alpha"], i+1);


      // loglike with current parameters

      old_loglike = cpp_ll_all(data, param, local_cases, list_custom_ll); // offset


      // proposal: swap case 'i' and its ancestor

      swapinfo = cpp_swap_cases(param, i+1);
      new_param["alpha"] = swapinfo["alpha"];
      new_param["t_inf"] = swapinfo["t_inf"];
      new_param["kappa"] = swapinfo["kappa"];


      // loglike with new parameters

      new_loglike = cpp_ll_all(data, new_param, local_cases, list_custom_ll);


      // acceptance term

      p_accept = exp(new_loglike - old_loglike);


      // acceptance: change param only if new values is accepted

      if (p_accept >= unif_rand()) { // accept new parameters
	param["alpha"] = new_param["alpha"];
	param["t_inf"] = new_param["t_inf"];
	param["kappa"] = new_param["kappa"];
      }
    }
  }

  return param;
}






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


// Movement of the number of generations on transmission chains ('kappa') is
// done for one ancestry at a time. As for infection times ('t_inf') we use a
// dumb, symmetric +/- 1 proposal. But because values are typically in a short
// range (e.g. [1-3]) we probably propose more dumb values here. We may
// eventually want to bounce back or use and correct for assymetric proposals.

// [[Rcpp::export(rng = true)]]
Rcpp::List cpp_move_kappa(Rcpp::List param, Rcpp::List data, Rcpp::List config,
			  Rcpp::RObject list_custom_ll = R_NilValue) {
  Rcpp::List new_param = clone(param);
  Rcpp::IntegerVector alpha = param["alpha"]; // pointer to param$alpha
  Rcpp::IntegerVector kappa = param["kappa"]; // pointer to param$kappa
  Rcpp::IntegerVector t_inf = param["t_inf"]; // pointer to param$t_inf
  Rcpp::IntegerVector new_kappa = new_param["kappa"];

  size_t N = static_cast<size_t>(data["N"]);
  size_t K = static_cast<size_t>(config["max_kappa"]);
  size_t jump;

  double old_loglike = 0.0, new_loglike = 0.0, p_accept = 0.0;

  for (size_t i = 0; i < N; i++) {

    // only non-NA ancestries are moved
    if (alpha[i] != NA_INTEGER) {

      // propose new kappa
      jump = (unif_rand() > 0.5) ? 1 : -1;
      new_kappa[i] = kappa[i] + jump;


      // only look into this move if new kappa is positive and smaller than the
      // maximum value; if not, remember to reset the value of new_kappa to that
      // of kappa, otherwise we implicitely accept stupid moves automatically

      if (new_kappa[i] < 1 || new_kappa[i] > K) {
	new_kappa[i] = kappa[i];
      } else {


	// loglike with current parameters
	old_loglike = cpp_ll_all(data, param, i+1, list_custom_ll);

	
	// loglike with new parameters
	new_loglike = cpp_ll_all(data, new_param, i+1, list_custom_ll);


	// acceptance term
	p_accept = exp(new_loglike - old_loglike);

	// acceptance: change param only if new values is accepted
	if (p_accept >= unif_rand()) { // accept new parameters
	  //	   Rprintf("\naccepting kappa:%d  (p: %f  old ll:  %f  new ll: %f",
	  // 		new_kappa[i], p_accept, old_loglike, new_loglike);
	  kappa[i] = new_kappa[i];
	} else {
	  new_kappa[i] = kappa[i];
	}
      }
    }

  }

  return param;
}