########################################################################################################################
# The TMB C code template :
# Mixed-effects model (MM) approach for estimating community dynamic parameters
# cf. MM approach of Tsai et al 2022 vs. IML approach of Engen et al. 2002
# CAUTION: The space ahead of the line is needed for C code template before submitting to TMB compiler
########################################################################################################################
library(TMB)
engen.cpp <- "// Non-linear mixed model specified through sparse design matrices.
#include <TMB.hpp>
template<class Type>
Type objective_function<Type>::operator() ()
{
  DATA_VECTOR(rho);         // Observations
  DATA_VECTOR(lag);         // Observations
  DATA_SPARSE_MATRIX(Z);    // Random effect design matrix
  PARAMETER_VECTOR(u_a);    // Random effects vector
  PARAMETER_VECTOR(u_i);    // Random effects vector
  PARAMETER_VECTOR(u_d);    // Random effects vector
  PARAMETER(logasymp);      // Fixed effects vector
  PARAMETER(loginit);       // Fixed effects vector
  PARAMETER(logdelta);      // Fixed effects vector
  PARAMETER(logsdu_a);      // Random effect standard deviations
  PARAMETER(logsdu_i);      // Random effect standard deviations
  PARAMETER(logsdu_d);      // Random effect standard deviations
  PARAMETER(logsd0);        // Measurement standard deviation
  // Distribution of random effect (u_a, u_i, u_d):
  Type ans = 0;
  ans -= dnorm(u_a, Type(0), exp(logsdu_a), true).sum();
  ans -= dnorm(u_i, Type(0), exp(logsdu_i), true).sum();
  ans -= dnorm(u_d, Type(0), exp(logsdu_d), true).sum();
  // Distribution of observations given random effects (y|u):
  vector<Type> Fasymp = exp(logasymp + Z*u_a);
  vector<Type> Fdelta = exp(logdelta + Z*u_d)*lag;
  vector<Type> Finit = exp(loginit + Z*u_i);
  vector<Type> y = (Finit-Fasymp)*exp(-Fdelta) + Fasymp;
  ans -= dnorm(rho, y, exp(logsd0), true).sum();
  return ans;
}"
writeLines(engen.cpp,con="ltmpdsr_engen.cpp")
compile("ltmpdsr_engen.cpp")
dyn.load(dynlib("ltmpdsr_engen"))
########################################################################################################################


TMB_MMestimate_fn <- function(){
  # Following R script to produce the mixed effects model estimates (estimating community dynamic parameters)
  load("dsrltmp2024_meta.RData")
  poilog_out <- poilog.mle.out 
  #data "poilog.mle.out" containing all MLE results of Poisson-lognormal fits across 39 reefs for all pairs of time difference (# see README file meta data)
  library(TMB) 
  library(Matrix)
  library(tidyverse)
  # loading required packages for TMB mixed model estimation 
  # Warning: TMB C code pre-compilation is required, which is a standard method please refer to TMB R package tutorial
  dyn.load(dynlib("ltmpdsr_engen")) 
  # the path of pre-compiled TMB code provided in "Rfunctions_sourcecode_dsrltmp2024.R"
  fit.data.all <- data.table::rbindlist(poilog_out)
  fit.data.all$reefname <- factor(as.character(fit.data.all$reefname))
  fit.data.all <- fit.data.all[-which(fit.data.all$reefname=="OPAL (2)" & fit.data.all$m==2022),]
  fit.data.all <- fit.data.all %>% mutate_at(c("mu1","mu2","sig1","sig2","rho","ll","conv","strue.x","strue.y"),as.numeric)
  #re-structure the data format as "fit.data.all" 
  Z <- model.matrix(~reefname-1,fit.data.all)
  Z <- as(Z,"TsparseMatrix")
  simdata <- list(rho=fit.data.all$rho,lag=fit.data.all$tlag,Z=Z)
  # pre-processing data structure as "simdata" for TMB estimation (values of rho and time lags)
  par <- list(u_a=rnorm(ncol(Z))*0.01,u_i=rnorm(ncol(Z))*0.01,u_d=rnorm(ncol(Z))*0.01,
              logasymp=-0.2,loginit=-0.05,logdelta=-0.1,
              logsdu_a=-3,logsdu_i=-5,logsdu_d=-3,logsd0=-2)
  # parameter initialization as "par" for TMB estimation, including fixed and random effects
  obj <- MakeADFun(data=simdata,parameters=par,random=c("u_a","u_i","u_d"),DLL="ltmpdsr_engen",hessian=TRUE,silent=TRUE)
  #  Use TMB's automatic differentiation "MakeADFun" to prepare for optimization
  opt <- suppressWarnings(nlminb(obj$par,obj$fn,obj$gr))
  #optimization procedure using R CaDENCE package (other optimization methods can be used alternatively)
  report <- sdreport(obj)
  # standard TMB report
  parfix <- report$par.fixed
  # obtain fixed effect parameters
  parrandom <- report$par.random
  # obtain random effect parameters
  asymp <- exp(parfix[1]+parrandom[1:39])
  # "rho_infinity" parameter derived from TMB for 39 reefs (see Table S1)
  init <- exp(parfix[2]+parrandom[40:78])
  # "rho_initial" parameter derived from TMB for 39 reefs (see Table S1)
  delta <- exp(parfix[3]+parrandom[79:117])
  # "delta" parameter derived from TMB for 39 reefs (see Table S1)
  Vr <- asymp
  # Proportion of variance in relative species abundance explained by persistent species differences
  Ve <- init-asymp
  # Proportion of variance in relative species abundance explained by environmental stochasticity
  Vd <- 1-init
  # Proportion of variance in relative species abundance explained by demographic/sampling stochasticity
  CDI2 <- Vr/Ve 
  # Community Determinism Index (CDI) defined in the main text equation 6 (demographic/sampling stochasticity excluded)
  reefname <- colnames(Z)
  return(data.frame(asymp,init,delta,Vr,Ve,Vd,CDI2,reefname))
}