// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
// There are three debug levels (o to 2).  0 will show now degugging output.
#define DEBUG 2

#if DEBUG == 1
#define DEBUG_OUTER
#endif

#if DEBUG == 2
#define DEBUG_INNER
#define DEBUG_OUTER
#endif

// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

// Macro for printing debug output
#define PRINTDBG(x) Rcout << __FUNCTION__ << "> " << #x << std::endl << x << std::endl
#define PRINTDBG_MSG(x) Rcout << __FUNCTION__ << "> " << x << std::endl

#include <Rcpp.h>
using namespace Rcpp;

// [[Rcpp::export()]]
NumericVector cpp_rvonmises(int n, double mu, double kappa) {
  /* Harry Southworth, January 2005.
   GENERATE RANDOM NUMBERS FROM A vonMises DISTRIBUTION.
   FOLLOWING THE ALGORITHM ON PAGE 43 OF MARDIA AND JUPP,
   DUE TO BEST AND FISHER.
   NOTE TAKEN OF ULRIC LUND'S IMPLEMENTATION IN S.
   THE ALGORITHM IS A REJECTION ALGORITHM AND RUNS
   SLOWLY IN S.
   Ported the C to R by Claudio Agostinelli, August 2006.

   Modified by James Campbell on 2023-06-22 for use in Rcpp code. Function written by James Campbell
   */
  
  if(kappa == 0){
    return(runif(n, -1*M_PI, M_PI));
  }

  // init output vector
  NumericVector x(n);

  
  double U1, U2, U3;
  double a, b, r, z, f, c;

  /* SET a, b AND r */
  a = 1 + sqrt(1+4* kappa * kappa);
  b = (a - sqrt(2*a))/(2* kappa);
  r = (1 + b*b)/(2*b);

  int i = 0;

  GetRNGstate();
  do{
    /* GENERATE U(0,1) RANDOM NUMBERS WITH THE R ROUTINE */
    U1 = unif_rand();
    z = cos(M_PI * U1);
    f = (1. + r * z)/(r + z);
    c = kappa * (r - f);

    U2 = unif_rand();
    if(c * (2 - c) - U2 > 0) {
      U3 = unif_rand();
      if (U3 > 0.50) x[i] = acos(f) + mu;
      else x[i] = -acos(f) + mu;
      i++;
    }
    else {
      if(log(c/U2) + 1 - c >= 0.) {
        U3 = unif_rand();
        if (U3 > 0.50) x[i] = acos(f) + mu;
        else x[i] = -acos(f) + mu;
        i++;
      }
    }

  } while(i < n);
  PutRNGstate();

  return(x);

}



enum TRANSFORM{NONE = 1, LOG = 2, COS = 3};
enum STEP_COL{X = 0, Y = 1, LEN = 2, BEARING = 3, TA = 4};


// [[Rcpp::export()]]
double angle(double x){
  if(x > 2*M_PI){
    return(fmod(x, 2*M_PI));
  }
  if(x < -2*M_PI){
    return(fmod(x, 2*M_PI) + 2*M_PI);
  }
  if(x < 0){
    return(x + 2*M_PI);
  }
  return(x);
}

NumericVector angle(NumericVector &x){
  
  for(int i = 0; i < x.size(); i++){
    x(i) = angle(x(i));
  }
  
  return(x);
  
}


// Calculate random step vector
// 
// Return a matrix where each row is a pseudo step. Columns 1:2 are the x/y
// coords, and columns 3:4 are the step length and bearing. Bearing is calculate
// from the vector difference between two step positions.
// 
// [[Rcpp::export()]]
void r_steps(
    NumericMatrix &new_steps, 
    NumericVector &current_step,
    double shape, double scale, double kappa){

  int n = new_steps.nrow();
  NumericMatrix::Column col_x = new_steps( _ , STEP_COL::X);
  NumericMatrix::Column col_y = new_steps( _ , STEP_COL::Y);
  NumericMatrix::Column col_step = new_steps( _ , STEP_COL::LEN);
  NumericMatrix::Column col_bear = new_steps( _ , STEP_COL::BEARING);
  NumericMatrix::Column col_ta = new_steps( _ , STEP_COL::TA);
  
  // step lengths
  col_step = rgamma(n, shape, scale);

  // bearing  
  NumericVector tmp;
  tmp = cpp_rvonmises(n, current_step(STEP_COL::BEARING), kappa);
  col_bear = angle(tmp);
  
  // turning angle
  tmp = tmp - current_step(STEP_COL::BEARING);
  col_ta = angle(tmp);

  // x/y position
  // Here, we're treating angles like geographic system, where north is 0
  // degrees, pointing upwards (+y) and east is clockwise (+x).
  col_x = -sin(-col_bear) * col_step + current_step[STEP_COL::X];
  col_y = cos(-col_bear) * col_step + current_step[STEP_COL::Y];
  
#ifdef DEBUG_INNER
  PRINTDBG(new_steps);
#endif
}
  
// Extract covariate data from RasterLayer
// [[Rcpp::export()]]
double extract_covariate(List &iSSF_raster, NumericVector &step){
  
  // Raster attributes
  NumericVector dims = iSSF_raster["dims"];
  NumericVector extent = iSSF_raster["extent"];
  bool directional = iSSF_raster["directional"];
  NumericMatrix values = iSSF_raster["values"];
  
  // Check if out of bounds
  if( step[0] < extent[0] | 
      step[0] > extent[1] |  
      step[1] < extent[2] | 
      step[1] > extent[3] ){
    return R_NaN;
  }
  
  // Get raster index to extract
  int icol = (int) floor(((step[0] - extent[0]) / (extent[1] - extent[0])) * (dims[1]));
  int irow = (int) floor(((extent[3] - step[1]) / (extent[3] - extent[2])) * (dims[0]));
  double out = values(irow, icol);
  
  // Check if covariate data is directional
  if(directional == true){
    // Convert covariate into bearing difference
    out = angle(angle(out) - angle(step[STEP_COL::BEARING]));
    if(out > M_PI){
      out -= 2*M_PI;
    }else if(out < -M_PI){
      out += 2*M_PI;
    }
    out = cos(out);
  }
  
  return out;
  
}


// Get varaible value for formula variable
double design_var(List &iSSF_raster_list, NumericVector &step, NumericVector &current_step, int ras_idx, int ras_trns, int ras_prev){
  
  List iSSF_raster;
  
  double var;
  
  // Extract variable value
  switch(ras_idx) {
  case -1: 
    // Extract step length
      var = step[2];
      break;
  case -2:
      // Extract turning angle
      var = step[4];
      break;
  default:
    // Extract raster value
    iSSF_raster = iSSF_raster_list[ras_idx];
    if(ras_prev){
      var = extract_covariate(iSSF_raster, current_step);
    }else{
      var = extract_covariate(iSSF_raster, step);
    }
  }
  
  // Apply transformations
  switch (ras_trns) {
  case TRANSFORM::LOG:
    var = log(var);
    break;
  case TRANSFORM::COS:
    var = cos(var);
    break;
  }
  
  // #ifdef DEBUG_INNER
  //   PRINTDBG(step);
  //   PRINTDBG(var);
  //   PRINTDBG(ras_idx);
  //   PRINTDBG(ras_trns);
  // #endif

  
  return(var);
}

// Get value for foruma linear term
double design_term(List &iSSF_raster_list, NumericVector &step, 
                   NumericVector &current_step,
                   IntegerVector &ras_idx, 
                   IntegerVector &ras_trns,
                   IntegerVector &ras_prev){
  
  int n_var = ras_idx.size();
  
  // Extract first variable from term
  double ret = design_var(iSSF_raster_list, step, current_step, ras_idx[0], ras_trns[0], ras_prev[0]);
  
  // Apply remaining interaction variables
  for(int i = 1; i < n_var; i++){
    ret *= design_var(iSSF_raster_list, step, current_step, ras_idx[i], ras_trns[i], ras_prev[i]);
  }
  
  return ret;
}

// Combine extracted raster covariates to make design matrix.
// This step is required for interaction variables.
NumericMatrix design_matrix(
    List &iSSF_raster_list,
    NumericMatrix &steps,
    NumericVector current_step,
    List &ras_idxs, List &ras_trnsfs, List &ras_prev){
  
  int n_term = ras_idxs.size();
  int n_steps = steps.nrow();
  
  NumericVector step(4);
  List term;
  
  NumericMatrix X(n_steps, n_term);
  IntegerVector idx, trns, prev;
  
  // loop steps
  for(int i_step = 0; i_step < n_steps; i_step++){
    
    step = steps(i_step, _);
   
    // Loop and extract linear terms 
    for(int i_term = 0; i_term < n_term; i_term++){
      idx = ras_idxs[i_term];
      trns = ras_trnsfs[i_term];
      prev = ras_prev[i_term];
      
      X(i_step, i_term) = design_term(iSSF_raster_list, step, current_step, idx, trns, prev);
    }
  }
  
#ifdef DEBUG_INNER
  PRINTDBG(X);
#endif
  
  return(X);
}


// Calculate step probability form covairate data
// [[Rcpp::export()]]
NumericVector step_probability_rast(NumericMatrix &X, NumericVector &beta_ras){
  
  int n_step = X.nrow();
  int n_param = beta_ras.size();
  NumericVector numerator(n_step); 
  NumericVector p(n_step);
  
  double tmp;
  
  // Loop steps
  for(int i_step = 0; i_step < n_step; i_step++){
    
    // Set init probability
    tmp = 0;
    
    // Loop beta parameters
    for(int i_par = 0; i_par < n_param; i_par ++){
      
      // Set probability to 0 if no covariate values
      if( R_isnancpp(X(i_step, i_par)) ){
        tmp = R_NegInf;
        break;
      }
      tmp += X(i_step, i_par) * beta_ras[i_par];
    }
    numerator(i_step) = exp(tmp);
  }
  
  double denominator = sum(numerator);
  if(denominator == 0){
    p.fill(0);
  }else{
    p = numerator / denominator;
  }
  

#ifdef DEBUG_INNER
  PRINTDBG(p);
#endif
  
  return(p);
  
}



//' Simulate an iSSF agent
//' 
//' This function will simulated the movement of an agent over a provided
//' landscape of covariate rasters.
//' 
//' @param position_start Coordinate the agent will start at.
//' @param bearing_start Initial bearing of the agent.
//' @param n_steps The number of steps to simulate per track.
//' @param len_track The number of tracks to simulate.
//' @param ras_lst A list of raster covariate data, returned from `iSSF_raster_list()`.
//' @param beta_par A vector holding the beta values for the model.
//' @param beta_ras_idx A list holding indices of the raster layers for each term in `beta_par`.
//' @param scale Scale parameter for gamma step length distribution. 
//' @param shape Shape parameter for gamma step length distribution. 
//' @param kappa kappa parameter for von Mises turning angle distribution.
//' 
// [[Rcpp::export()]]
List cpp_iSSF_simulation(
    NumericVector position_start, double bearing_start, 
    int n_steps, int len_track, 
    List iSSF_raster_list, 
    NumericVector coef,
    List ras_idx, List ras_trnsfs, List ras_prev,
    double scale, double shape, double kappa
    ){
 
  // Get number of raster layers
  int n_rasters = iSSF_raster_list.size();
  
  // Get number of model parameters
  int n_coef = coef.size();
  
  // Init agent data objects
  NumericMatrix positions(len_track + 1, 2); 
  positions(0,_) = position_start;
  
  NumericVector bearing(len_track + 1); 
  bearing(0) = bearing_start;
  
  NumericVector step_length(len_track + 1); 
  step_length(0) = 0; 

  NumericVector turning_angle(len_track + 1); 
  turning_angle(0) = 0; 
  
  NumericVector current_step(5);
  current_step[seq(0,1)] = position_start;
  current_step[2] = 0;
  current_step[3] = bearing_start;
  current_step[4] = 0;
  
  // Init step selection probability vector  
  double p_vals[n_steps];

#ifdef DEBUG_OUTER
  PRINTDBG(position_start);
  PRINTDBG(bearing_start);
  PRINTDBG(n_steps);
  PRINTDBG(len_track);
  PRINTDBG(scale);
  PRINTDBG(shape);
  PRINTDBG(kappa);
#endif  

  // temp variables for holding agent position and possible new steps
  NumericMatrix new_steps(n_steps, 5);

  // Matrix holding raster covariates (columns) for each possible new step (rows)
  NumericMatrix ras_vals(n_steps, n_rasters);
  // Design matrix
  NumericMatrix X(n_steps, n_rasters);
  
  NumericVector p(n_steps);  
  
  // rmultinom requires a pointer to store the return value.
  int drawn[n_steps];
  int select;

  // Make objects for holding list elements 
  List ras_lst_element;
  IntegerVector beta_ras_idx_element; 
  
  // ============= loop time steps and simulate agent movement =================
  for(int i = 1; i <= len_track; i++){

#ifdef DEBUG_OUTER
    PRINTDBG(i);
#endif


#ifdef DEBUG_OUTER
    PRINTDBG_MSG("Make new steps");
#endif
    // ------ Populate new_step matrix
    // Make new random steps
    r_steps(new_steps, current_step, shape, scale, kappa);

#ifdef DEBUG_OUTER
    PRINTDBG_MSG("Make design matrix");
#endif
    // ------ Get linear design matrix
    // i.e. combine values for interacting covariate layers.
    X = design_matrix(iSSF_raster_list, new_steps, current_step, ras_idx ,ras_trnsfs, ras_prev);

#ifdef DEBUG_OUTER
    PRINTDBG_MSG("Get new step probabilities");
#endif
    // ------ Calculate step probabilities: Raster covariates
    p = step_probability_rast(X, coef);

    // If all steps are out of bounds, stay still, skip to next step
    if(sum(p) == 0){
      positions(i, 0) = positions(i - 1, 0);
      bearing(i) = bearing(i-1);
      turning_angle(i) = turning_angle(i-1);
      step_length(i) = R_NaN;
      continue;
    }

    // Draw new step
    // rmunlinom docs:
    // https://dirk.eddelbuettel.com/code/rcpp/html/Rmath_8h_source.html#l00102
    // https://stackoverflow.com/questions/63368639/how-to-use-multinomial-function-in-c-with-rmath-h
    rmultinom(1, p.begin(), n_steps, drawn);
    for(int j = 0; j < n_steps; j++){
      if(drawn[j] != 0){
        select = j;
        break;
      }
    }
    
    // update current step
    current_step = new_steps(select, _);
    
    // write output
    positions(i, 0) = current_step[0];
    positions(i, 1) = current_step[1];
    step_length[i] = current_step[2];
    
#ifdef DEBUG_OUTER
    PRINTDBG(select);
#endif
    
    
  }
  
  // ------ Return an agent
  List agent = List::create(
    _["track"] = positions,
    _["step_length"] = step_length,
    _["bearing"] = bearing,
    _["turning_angle"] = turning_angle);
  
  return(agent);
  
}