#' ---
#' title: "Vignette for simulation analysis"
#' author: " "
#' date: "`r format(Sys.time(), '%d %B, %Y')`"
#' output: html_document
#' ---
#'
#'
#' ## Setup
## ---- warning=FALSE, message=FALSE--------------------------
library(lubridate)
library(terra)
library(raster)
library(tidyverse)
library(tidyterra)
library(conflicted)
library(ggplot2)
library(glmmTMB)
library(amt) # dev version needed: remotes::install_github("jmsigner/amt")
library(cowplot)
library(MASS)
library(mixedSSA) # install from devtools::install_github("smthfrmn/mixedSSA")

conflict_prefer("select", "dplyr")
conflict_prefer("filter", "dplyr")

set.seed(11102023)
n <- 30 ### Number of individuals to simulate 

#' ## Introduction
#'
#' We simulate trajectories following **Integrated step-selection analysis (ISSA)** (Equation 1) for a population with individual-specific, 
#' known parameter values to demonstrate the application of mixed ISSAs that allow for individual variability 
#' in movement parameters (Figure 1). Individual tracks were simulated using the `amt` package. 
#' We also compare the estimates from the mixed ISSA to those obtained by fitting a separate model to true 
#' individual-specific parameters from each individual. This comparison highlights the benefits of using random effects in ISSA; namely; the 
#' ability to borrow information across like individuals, which leads to improved estimates of individual-specific parameters.
#' 
#' ## Raster Generation
#' We generated a categorical raster with two classes (Habitat A, Habitat B) where the individual selects
#' for the habitat B and moves differently in the two habitats (Section 3.1). 

tc2 <- raster(nrows=1000, ncols=1000, xmn=-500, xmx=500, ymn=-500,ymx=500)  
tc2[] <- runif(1000000, 0,50) ## assigning simulated values to the raster
# producing a new raster with moving window of size 11x13 and calculating mean values for the neighborhood
tc2 <- raster::focal(tc2, w = matrix(1, 11, 13), mean) 
tc2 <- raster::focal(tc2, w = matrix(1, 11, 13), mean)
# we did the focal operation twice to produce a smoothed out edges of the interspersed habitats
reclass_df  <-matrix(c(12,24.5, 1, 24.5,38, 0), ncol=3, byrow=T)
tc_reclass <- raster::reclassify(tc2, reclass_df)

r <- rast(tc_reclass)
names(r) <- "x1"
r <- as.factor(r)
cls <- data.frame(id=c(0,1), x1=c("Habitat A", "Habitat B"))
levels(r) <- cls
plot(r)

#+ eval =FALSE
#### Initializing
sig <-matrix(c(0.25,0.15,0.15, 0.15, 0.2, 0.15, 0.15,0.15, 0.25), ncol =3, byrow = T)
param.value <-data.frame(id = NA, shape = NA, scale = NA, kappa = NA)
ssf_sim <-list()

#' ## Simulating tracks using the amt functions
#' The simulation of movement tracks in *amt* requires three steps: 
#' 1. Specifying parameters governing the movements for the individual using the make_issf_model function   
#' 2. Specifying a redistribution kernel for generating movements based on an initial location and bearing   
#' 3. Generating paths from the simulated redistribution kernel specified for a user-provided number of steps    
#'   
#' **Warning**: The simulation takes a while to run (a few hours to days) and requires amt version 0.2.2. 
#' Please check the version and compatibility before the simulation run.

#+ eval =FALSE
for(i in 1:n){
  start <- make_start(c(runif(1,-20,20), runif(1,-20,20)), ta_ = 0)
  param.value[i,1] <- i
  param.value[i,2:4] <- mvrnorm(n=1, mu=c(3,1,4), Sigma = sig) 
  ## simulating initial parameters using multivariate normal distribution
  
  m <- make_issf_model(coefs = c(x1_end = 0.05,  
                                 "sl_" = 0, "log(sl_)" = 0, "cos(ta_)" = 0,
                                 "sl_:x1_start" = 0.05, "log(sl_):x1_start" = 0.05,"cos(ta_):x1_start" = 0.05), 
                       sl = make_gamma_distr(shape = param.value[i,2], 
                                             scale = param.value[i,3]
                       ), 
                       ta = make_vonmises_distr(kappa = param.value[i,4]
                       ))
  
  rdk <- redistribution_kernel(m, start = start, map = r, as.rast= FALSE,
                               landscape= "continuous")
  
  ssf_sim[[i]] <- simulate_path(rdk, n.steps = floor(rnorm(1,2000,250))) ### simulate individual trajectories
  ssf_sim[[i]]$id <-i
}
#+ echo= FALSE
param.value <-read.csv("Param_all_1510.csv")[,1:4]
#+ echo= TRUE
#+ Initial movement parameter values used for the simulation
head(param.value)

#+ eval =FALSE
ssf_simdat <-do.call(rbind, ssf_sim) ### binding individual tracks into one dataset

#'  Making tracks from the simulated trajectories:
#+ eval =FALSE
ssf_simtr <- ssf_simdat |> make_track(x_, y_, t_ , id=id)
#+ echo= FALSE
ssfdat_sim_track_ind <-readRDS("Indv_simdata_final.rds")
ssf_simtr <- ssfdat_sim_track_ind |> filter(case_ == TRUE)

#' Visualization of the tracks:
#+ eval = TRUE, warning = FALSE, message = FALSE
ssf_simtr_sample <- ssf_simtr |> filter(id %in% c(1:4)) 

ggplot() + geom_spatraster(data= r, show.legend = FALSE)  +
  geom_path(data= ssf_simtr_sample, aes(x1_, y1_),  col= "black")+
  facet_wrap(~ id, nrow =2)+ scale_fill_terrain_d() + theme_bw()


#' ## Individual Model fitting
#' ### Generating random steps from the tracks
#+ eval =FALSE
ssfdat_sim_track_ind <- ssf_simtr %>% nest(data = -id) %>%
  mutate( data = map(data,  ~ .x %>% 
                       steps() %>%
                       random_steps() %>%
                       extract_covariates(r, where = "both"))) %>%
  unnest(cols = data)


#+ echo=TRUE
glimpse(ssfdat_sim_track_ind)

### Fitting Individual model 
ssffits <- ssfdat_sim_track_ind %>% nest(data=-id) %>% 
  mutate(mod = map(data, function(x) (try(fit_issf(case_ ~ sl_ + log(sl_)+ cos(ta_)+ x1_end   + 
                                                     sl_:x1_start + log(sl_):x1_start+ cos(ta_):x1_start+
                                                     strata(step_id_), data = x)))))


### Updating individual model
#' We estimate the tentative individual movement parameters for each individual and then 
#' update the movement parameters using the *update_gamma* and *update_vonmises* functions of the *amt* package. 
#' In this vignette, we will be focusing on the movement parameters when the individual is in the reference habitat
#' class (Habitat A).
for(i in 1:n){
  
  indv_dat <- subset(ssfdat_sim_track_ind, ssfdat_sim_track_ind$id == i)
  sl_fit <- amt::fit_distr(indv_dat$sl_[indv_dat$case_=="TRUE"], "gamma", na.rm = T)
  ta_fit <- amt::fit_distr(indv_dat$ta_[indv_dat$case_=="TRUE"], "vonmises", na.rm = T)
  
  tab <- broom::tidy(ssffits$mod[[i]]$model)
  
  param.value[i, 5:6] <- as.data.frame(update_gamma(dist= sl_fit, 
                                                    beta_sl = tab$estimate[tab$term == "sl_"] ,
                                                    beta_log_sl = tab$estimate[tab$term == "log(sl_)"])$params)[1,]
  
  
  param.value[i, 7:8] <- as.data.frame(update_vonmises(dist= ta_fit, 
                                                       beta_cos_ta = tab$estimate[tab$term == "cos(ta_)"])$params)[1,]
  
}
names(param.value)[5:8] <-c("shape.indv", "scale.indv", "kappa.indv", "mu.indv") 

head(param.value)

#' ## Mixed-effect Model fitting
#' ### Generating random steps from the tracks with one tentative distribution
#' The step to generate random steps is similar to the individual model but we are
#' using one tentative distribution to generate the random steps for all individuals.
#+ eval = FALSE
ssfdat_sim_track <- ssf_simtr %>% nest(data = -id) %>%
  mutate( data = map(data,  ~ .x %>% 
                       steps())) %>%
  unnest(cols = data) %>%
  random_steps() %>%
  extract_covariates(r, where = "both") 

ssfdat_sim_track$step_id <- with(ssfdat_sim_track, paste(id, step_id_, sep = "_"))
#+ echo= FALSE
ssfdat_sim_track <-readRDS("Mixed_simdata_final.rds")
#+ echo=TRUE
glimpse(ssfdat_sim_track)
#saveRDS(ssfdat_sim_track, "Mixed_simdata_final.rds")

#' Fitting the mixed-effect model using *glmmTMB* package following Muff et al. (2020):    
#' The code repository can be found at https://conservancy.umn.edu/handle/11299/204737
#+ eval =FALSE

ssf_sim_model <- glmmTMB(case_ ~    sl_ + log(sl_)+ cos(ta_)+ x1_end   +  
                           sl_:x1_start + log(sl_):x1_start + cos(ta_):x1_start +
                           (1|step_id)+ (0+ sl_ +(log(sl_)) + cos(ta_)|id) +(0+ x1_end|id), 
                         REML = TRUE, family=poisson(), data =ssfdat_sim_track,  doFit=FALSE)

## we have to manually fix the variance of the intercept first and then fix the 
## standard deviation of the first random term, which is the (1|step_id) component 
## in the above model equation that is fixed as 1e3 here.
ssf_sim_model$parameters$theta[1] <- log(1e3)

## We need to tell glmmTMB not to change the variance by setting it to NA:

nvarparm7 <- length(ssf_sim_model$parameters$theta)
ssf_sim_model$mapArg <- list(theta=factor(c(NA, 1:(nvarparm7-1))))
## Then fit the model and look at the results:
snapper.ssf_sim<- glmmTMB:::fitTMB(ssf_sim_model)
#saveRDS(snapper.ssf_sim, "population_model_final.rds")

#+ echo= FALSE
snapper.ssf_sim <- readRDS("population_model_final.rds")
#+ echo=TRUE
summary(snapper.ssf_sim)

#' After fitting the mixed-effect model using the *glmmTMB* package, we get the model output
#' into the *mixedSSA* package for updating the movement parameters. First, we update the parameters 
#' in the step-length distribution and assume they follow a *gamma* distribution.

#+ warning = FALSE, message = FALSE
updated_sl_parameters <- mixedSSA::update_dist(model = snapper.ssf_sim,
                                            dist_name = "gamma", # name of the step-length distribution assumed
                                            beta_sl =  "sl_", # the name of the step lengths coefficient in our model
                                            beta_log_sl = "log(sl_)", # the name of the log(sl) coefficient in our model
                                            interaction_var_name = "x1_start", # name of the habitat variable layer
                                            random_effects_var_name = "id", # for individual update
                                            quantiles = FALSE # for continuous interaction variables
)

# Extract the updated movement parameters for the reference habitat class and 
# not include the first row of tentative movement parameters
param.value[,9:10] <- updated_sl_parameters@updated_parameters[-1,5:6] 

#' Then, we update the parameters in the turn-angle distribution assuming a *von Mises* distribution.
#+ warning = FALSE, message = FALSE
updated_ta_parameters <- mixedSSA::update_dist(model = snapper.ssf_sim,
                                               dist_name = "vonmises", # name of the turn-angle distribution assumed
                                               beta_cos_ta =  "cos(ta_)", # the name of the cos(ta) coefficient in our model
                                               interaction_var_name = "x1_start", # name of the habitat variable layer
                                               random_effects_var_name = "id", # for individual update
                                               quantiles = FALSE # for continuous interaction variables
)

# Extract the updated movement parameters for the reference habitat class and 
# not include the first row of tentative movement parameters
param.value[,11:12] <- updated_ta_parameters@updated_parameters[-1,4:5] 
names(param.value)[9:12] <-c("shape.all", "scale.all", "kappa.all", "mu.all") 


#+ echo = TRUE
head(param.value)

#' ## Plotting for model comparison
#' 
#' We calculated the deviations between each estimated parameter and the parameter value used to simulate the data.
#' We compared the deviation of individual and mixed-effect model to evaluate their relative performance.

param.value1 <- param.value |> mutate(indv_shape_bias = shape - shape.indv,
                                      mixed_shape_bias = shape - shape.all,
                                      indv_scale_bias = scale - scale.indv,
                                      mixed_scale_bias = scale - scale.all,
                                      indv_kappa_bias = kappa - kappa.indv,
                                      mixed_kappa_bias = kappa - kappa.all) 


shp_bias <- param.value1 |> 
  pivot_longer(cols = c(indv_shape_bias, mixed_shape_bias), names_to = "shapediff", values_to = "shape_value")|>
  ggplot()+ geom_boxplot(aes(y= shape_value, x = shapediff, fill = shapediff), outlier.shape = NA)+
  geom_jitter(aes(y= shape_value, x = shapediff), size = 0.6)+
  scale_x_discrete(labels=c('Individual\nModel', 'Mixed-effect\nModel'))+
  scale_fill_manual(breaks = c("indv_shape_bias", "mixed_shape_bias"),
                    values = c("indv_shape_bias" = "#ab1489",
                               "mixed_shape_bias"= "#69b3a2"))+
  labs(y= "Shape Parameter deviation", x= "Model type")+  theme_bw()+
  theme(legend.position = "none", 
        axis.title = element_text(size=16)) 

scl_bias <- param.value1 |> 
  pivot_longer(cols = c(indv_scale_bias, mixed_scale_bias), names_to = "scalediff", values_to = "scale_value")|>
  ggplot()+ geom_boxplot(aes(y= scale_value, x = scalediff, fill = scalediff), outlier.shape = NA)+
  geom_jitter(aes(y= scale_value, x = scalediff), size = 0.6)+
  scale_x_discrete(labels=c('Individual\nModel', 'Mixed-effect\nModel'))+
  scale_fill_manual(breaks = c("indv_scale_bias", "mixed_scale_bias"),
                    values = c("indv_scale_bias" = "#ab1489",
                               "mixed_scale_bias"= "#69b3a2"))+
  labs(y= "Scale Parameter deviation", x= "Model type")+ theme_bw()+ 
  theme(legend.position = "none", 
        axis.title = element_text(size=16))

kap_bias <- param.value1 |>
  pivot_longer(cols = c(indv_kappa_bias, mixed_kappa_bias), names_to = "kappadiff", values_to = "kappa_value") |>
  ggplot()+ geom_boxplot(aes(y= kappa_value, x = kappadiff, fill = kappadiff), outlier.shape = NA)+
  geom_jitter(aes(y= kappa_value, x = kappadiff), size = 0.6)+
  scale_x_discrete(labels=c('Individual\nModel', 'Mixed-effect\nModel'))+
  scale_fill_manual(breaks = c("indv_kappa_bias", "mixed_kappa_bias"),
                    values = c("indv_kappa_bias" = "#ab1489",
                               "mixed_kappa_bias"= "#69b3a2"))+
  labs(y= "Kappa Parameter deviation", x= "Model type")+ theme_bw()+ 
  theme(legend.position = "none", 
        axis.title = element_text(size=16))

plot_grid(shp_bias, scl_bias, kap_bias, ncol=3, labels = c("A", "B", "C"))

sessionInfo()