# Reversible jump MCMC for correlated, biased random walk mixture model with c <-  3 centres of attraction and an exploratory state

set.seed(1)   # Set seed
#------------------------------------------------------------------------------#

# Specify model inputs
ncentres <- 3              # Number of centres of attraction (c)
nexplore <- 2              # Number of exploratory states (h)

iter <- 20000              # Number of iterations of the Markov chain
thin <- 1                  # Number of iterations for thinning of the Markov chain.  Must be a factor of iter.
adapt <- 20000             # Number of iterations for pilot tuning
bin <- 500                 # Number of iterations between each pilot tune and calculation of MH acceptance rates
taccept <- 0.3             # Target acceptance rate for Metropolis-Hastings updates
txyaccept <- 0.2           # Target acceptance rate for (X_t,Y_t) single MH updates
tcentreaccept <- 0.05        # Target acceptance rate for centre of attraction (X_z,Y_z) MH updates


sigma_scale <- 0.001       # Scale parameter for the inverse gamma prior on sigma2_x and sigma2_y
sigma_shape <- 0.001       # Shape parameter for the inverse gamma prior on sigma2_x and sigma2_y
tau_scale <- 2             # Scale parameter for the inverse gamma prior on tau
tau_shape <- 3             # Shape parameter for the inverse gamma prior on tau
b_upper <- 30              # Upper limit for Unif(0,bupper) prior on Weibull shape parameter (b)

#------------------------------------------------------------------------------#

# Load data for grey seal 908.  Easting coordinates in UTM zone 31 have been re-scaled to the distance eastward relative to the central meridian of UTM zone 30.
  # 'observed_data' are the GPS locations and times
  # 'complete_data' are the complete dataset (x_ti, y_ti), including missing observations where i=0 (indicated by 0) and times
  # 'predicted_times' are the times for predicted locations
  # 'ssidx' is a vector for looping over the observed location (x_ti, y_ti) within each time step t
  # 'j_ti' is a vector indicating the proportion of each interval (t-1,t) at which the ith observation (x_ti,y_ti) was obtained.  For any interval where i=0, j_ti <-  -1.
  # 'nogo1x' are the X coordinates of vertices for a polygon indicating inland areas
  # 'nogo1y' are the Y coordinates of vertices for a polygon indicating inland areas
  # 'nogo2x' are the X coordinates of vertices for a polygon indicating inland areas
  # 'nogo2y' are the Y coordinates of vertices for a polygon indicating inland areas
  # 'nogo3x' are the X coordinates of vertices for a polygon indicating inland areas
  # 'nogo3y' are the Y coordinates of vertices for a polygon indicating inland areas
  # 'nogo4x' are the X coordinates of vertices for a polygon indicating inland areas
  # 'nogo4y' are the Y coordinates of vertices for a polygon indicating inland areas
  # 'nvt' is the number of vertcies for the nogo polygons
  
load("seal908_data.RData")
N <- length(predicted_times)                       # Number of predicted locations (T+1)
x_ti <- complete_data[,"x_ti"]                     # Observed locations, x_ti, of the X coordinate at time t, i.  These include missing locations where i=0 (indicated by x_ti <-  0) 
y_ti <- complete_data[,"y_ti"]                     # Observed locations, y_ti, of the Y coordinate at time t, i.  These include missing locations where i=0 (indicated by y_ti <-  0) 
maxdist <- max(dist(cbind(observed_data[,"x"],observed_data[,"y"]),method="euclidean"))            # Calculate maximum distance between observed locations.  This is used to scale distance within tanh() function for the strength of bias to centres of attraction.            
maxspeed <- 2                                      # Maximum speed of animal (meters per second)
time_interval <- 120                               # Time interval between predicted locations (minutes)
maxstep <- maxspeed*time_interval*60                   # Maximum step length based on maxspeed and time interval (meters)

#------------------------------------------------------------------------------#

# Declare parameters and latent variables
eta_z <- numeric((iter/thin+1)*ncentres)                    # Correlation between consecutive movement directions in centre of attraction state z
m_z <- numeric((iter/thin+1)*ncentres)                      # Intercept term for strength of bias (on logit scale)
r_z <- numeric((iter/thin+1)*ncentres)                      # Strength of bias as a function of distance to centre of attraction z (on logit scale)
q_z <- numeric((iter/thin+1)*ncentres)                      # Strength of bias as a function of distance^2 to centre of attraction z (on logit scale)
X_z <- numeric((iter/thin+1)*ncentres)                      # Location of the X coordinate for centre of attraction z
Y_z <- numeric((iter/thin+1)*ncentres)                      # Location of the Y coordinate for centre of attraction z
a <- numeric((iter/thin+1)*(ncentres+nexplore))             # Scale parameter of the Weibull distribution for step length (for centres of attraction, when delta_t > d_z)  
b <- numeric((iter/thin+1)*(ncentres+nexplore))             # Shape parameter of the Weibull distriubtion for step length (for centres of attraction, when delta_t > d_z)    
a2 <- numeric((iter/thin+1)*ncentres)                       # Scale parameter of the Weibull distribution for step length (when delta_t < d_z)  
b2 <- numeric((iter/thin+1)*ncentres)                       # Shape parameter of the Weibull distriubtion for step length (when delta_t < d_z)  
d_z <- numeric((iter/thin+1)*ncentres)                      # Change-point distance of Weibull distribution parameters for centre of attraction states
upsilon <- numeric((iter/thin+1)*nexplore)                  # Correlation between consecutive movement directions in an exploratory state
phi_0 <- numeric(iter/thin+1)                               # Initial movement direction
tau <- numeric(iter/thin+1)                                 # Variance of the Normal(0,tau) prior distributions for r_z and q_z
sigma2_x <- numeric(iter/thin+1)                            # Variance of the N(0,sigma2_x) observation error term in the x direction
sigma2_y <- numeric(iter/thin+1)                            # Variance of the N(0,sigma2_y) observation error term in the y direction
X_t <- numeric((iter/thin+1)*N)                             # Predicted location of the X coordinate at time t
Y_t <- numeric((iter/thin+1)*N)                             # Predicted location of the Y coordinate at time t
psi <- numeric((iter/thin+1)*(ncentres+nexplore)^2)                # State transition probabilities
z_t <- integer((iter/thin+1)*N)                             # State assignment for each time step t, including initial state z_0

centre_loc <- numeric((iter/thin+1)*ncentres)               #Vector indicating which predicted locations (X_t, Y_t) are the current centres of attraction
model_selection_parameters <- ncentres*2+nexplore           #Number of parameters included in model selection and multimodel inference
Model <- integer((iter/thin+1)*model_selection_parameters)  #Vector indicating the model structure; '1' indicates a (state-specific) parameter currently in the model from the full set ordered by eta_z, q_z, and upsilon (for z <-  1,...,c+h).  For example, for c=3, 01110010 indicates eta_2, eta_3, q_1, and upsilon_4 are currently in the model 

#------------------------------------------------------------------------------#

# Load initial values for the Markov chain
load("initial_values.RData")

# Assign initial values to parameters. These should be mofified if c < 3 or c > 3.
eta_z[1:ncentres] <- initial_values$eta_z
m_z[1:ncentres] <- initial_values$m_z
r_z[1:ncentres] <-   initial_values$r_z
q_z[1:ncentres] <- initial_values$q_z
X_z[1:ncentres] <- initial_values$X_z
Y_z[1:ncentres] <- initial_values$Y_z
a[1:(ncentres+nexplore)] <- initial_values$a
b[1:(ncentres+nexplore)] <- initial_values$b  
a2[1:ncentres] <- initial_values$a2
b2[1:ncentres] <- initial_values$b2
d_z[1:ncentres] <- initial_values$d_z
upsilon[1:nexplore] <- initial_values$upsilon
phi_0[1] <- initial_values$phi_0
tau[1] <- initial_values$tau
sigma2_x[1] <- initial_values$sigma2_x
sigma2_y[1] <- initial_values$sigma2_y
X_t[1:N] <- initial_values$X_t
Y_t[1:N] <- initial_values$Y_t
psi[1:(ncentres+nexplore)^2] <- initial_values$psi
z_t[1:N] <- initial_values$z_t

centre_loc[1:ncentres] <- initial_values$centre_loc
Model[1:model_selection_parameters] <- initial_values$Model

#------------------------------------------------------------------------------#

# Other variables
firstparms <- ncentres*4+ncentres*5+2*(ncentres+nexplore)+nexplore+2+2   # Number of parameters updated using Metropolis-Hastings (MH) steps, excluding predicted locations (X_t, Y_t) 
gs <- initial_values$gs                              # Tuning variables for (adaptive) proposal distributions.  These should be mofified if c < 3 or c > 3.
arate <- rep(0,firstparms+2*N)                       # Vector to store acceptance rate of parameters
loglike <- numeric(iter/thin+1)                      # Vector to store values for the log-likelihood function

#------------------------------------------------------------------------------#

# system() function is used to compile C code and create dynamic-link library file "RJMCMC.dll"
#system(paste("Rcmd SHLIB RJMCMC.c"))

# MCMC iteration count will print in R console if print_iteration_count =  1.  If print_iteration_count =  0, then iteration count will not print.
print_iteration_count <- 1 

# Load the compiled .dll
dyn.load("RJMCMC.dll")

# Using .C interface, draw from the posterior distribution using the reversible jump MCMC algorithm
posterior <-  .C("MCMCloop",as.numeric(eta_z),as.numeric(m_z),as.numeric(r_z),as.numeric(q_z),as.numeric(X_z),as.numeric(Y_z),as.numeric(a),as.numeric(b),as.numeric(a2),as.numeric(b2),as.numeric(d_z),as.numeric(upsilon),as.numeric(phi_0),as.numeric(tau),as.numeric(sigma2_x),as.numeric(sigma2_y),
                as.numeric(X_t),as.numeric(Y_t), as.numeric(psi),as.integer(z_t),as.integer(centre_loc),as.integer(Model),
                as.numeric(sigma_scale), as.numeric(sigma_shape),as.numeric(tau_scale),as.numeric(tau_shape),as.numeric(b_upper),
                as.numeric(loglike),as.numeric(gs),as.numeric(arate),
                as.numeric(x_ti),as.numeric(y_ti),as.integer(N),as.integer(firstparms),as.integer(ncentres),as.integer(nexplore),as.integer(ssidx),as.numeric(j_ti),
                as.integer(nvt),as.numeric(nogo1x),as.numeric(nogo1y),as.numeric(nogo2x),as.numeric(nogo2y),as.numeric(nogo3x),as.numeric(nogo3y),as.numeric(nogo4x),as.numeric(nogo4y),as.numeric(maxdist),as.numeric(maxstep),
                as.integer(iter),as.integer(thin),as.integer(adapt),as.integer(bin),as.numeric(taccept),as.numeric(txyaccept),as.numeric(tcentreaccept),as.integer(print_iteration_count))

# Unload the compiled .dll
dyn.unload("RJMCMC.dll")

# Label the elements of posterior
names(posterior)=c("eta_z","m_z","r_z","q_z","X_z","Y_z","a","b","a2","b2","d_z","upsilon","phi_0","tau","sigma2_x","sigma2_y","X_t","Y_t","psi","z_t","centre_loc","Model","sigma_scale","sigma_shape","tau_scale","tau_shape","bupper","loglike","gs","arate","x_ti","y_ti","N","firstparms","ncentres","nexplore","ssidx","j_ti","nvt","nogo1x","nogo1y","nogo2x","nogo2y","nogo3x","nogo3y","nogo4x","nogo4y","maxdist","maxstep","iter","thin","ada","bin","taccept","txyaccept","tcentreaccept","print_iteration_count")

#------------------------------------------------------------------------------#

# Process posterior vector output
eta_z <-  matrix(posterior$eta_z,ncol = ncentres,byrow = T)
r_z <-  matrix(posterior$r_z,ncol = ncentres,byrow = T)
q_z <-  matrix(posterior$q_z,ncol = ncentres,byrow = T)
m_z <-  matrix(posterior$m_z,ncol = ncentres,byrow = T)
X_z <-  matrix(posterior$X_z,ncol = ncentres,byrow = T)
Y_z <-  matrix(posterior$Y_z,ncol = ncentres,byrow = T)
a <-  matrix(posterior$a,ncol = ncentres+nexplore,byrow = T)
b <-  matrix(posterior$b,ncol = ncentres+nexplore,byrow = T)
a2 <-  matrix(posterior$a2,ncol = ncentres,byrow = T)
b2 <-  matrix(posterior$b2,ncol = ncentres,byrow = T)
d_z <-  matrix(posterior$d_z,ncol = ncentres,byrow = T)
upsilon <-  matrix(posterior$upsilon,ncol = max(1,nexplore),byrow = T)
phi_0 <-  posterior$phi_0
tau <-  posterior$tau
sigma2_x <-  posterior$sigma2_x
sigma2_y <-  posterior$sigma2_y
X_t <-  matrix(posterior$X_t,ncol = N,byrow = T)
Y_t <-  matrix(posterior$Y_t,ncol = N,byrow = T)
psi <-  matrix(posterior$psi,ncol = (ncentres+nexplore)*(ncentres+nexplore),byrow = T)
z_t <-  matrix(posterior$z_t,ncol = N,byrow = T)
centre_loc <-  matrix(posterior$centre_loc,ncol = ncentres,byrow = T)
Model <-  matrix(posterior$Model,ncol = model_selection_parameters,byrow = T)

# Model Selection Results
  library(gtools)
  PMW <-  matrix(nrow=(2^model_selection_parameters),ncol=model_selection_parameters)
  PMW <-  permutations(2,model_selection_parameters,repeats.allowed=TRUE)
  PMW[PMW==1] <-  0
  PMW[PMW==2] <-  1

  Modelsums <- numeric(2^model_selection_parameters)
  for(g in 1:(2^model_selection_parameters))	{
    tempPMW <-  matrix(PMW[g,],ncol=ncol(Model),nrow=nrow(Model),byrow=T)
    md <-  (tempPMW==Model)
    Modelsums[g] <-  sum(rowSums(md)==model_selection_parameters)
    rm(tempPMW)
  }

  ModelResults <-  cbind(Modelsums/nrow(Model),PMW)
  ModelResultsTable <-  ModelResults[which(ModelResults[,1]>0),]
  ModelResultsTable <-  ModelResultsTable[order(ModelResultsTable[,1]),]
  parmlab=character(2*ncentres+nexplore)
  for(i in 1:ncentres) {
    parmlab[i]=paste("eta_",i,sep="")
    parmlab[i+ncentres]=paste("q_",i,sep="")
  }
  if(nexplore>0)  {
    for(i in 1:nexplore) {
      parmlab[i+2*ncentres]=paste("upsilon_",i+ncentres,sep="")
    }
  }
  colnames(ModelResultsTable) <-  c("Weight",parmlab)
  print(ModelResultsTable)