setwd("")

counts <- as.matrix(read.delim("counts.txt"))
colnames(counts) <- NULL
rownames(counts) <- NULL

library(nimble)
library(truncnorm)

source("nimble_model.R")

T0 <- rep(NA,dim(counts)[1])
Tf <- rep(NA,dim(counts)[1])
for (i in 1:dim(counts)[1]){
  T0[i] <- min(which(is.na(counts[i,]) == FALSE))
  Tf[i] <- max(which(is.na(counts[i,]) == FALSE))
}
names(T0) <- names(Tf) <- rownames(counts)

# DATA
win.data <- list(y = counts)
constants <- list(nSite = nrow(counts),nYear = ncol(counts),
                   T0 = T0, Tf = Tf)

# Parameters to save
param.nb <- c("N","rho","eta","sigma.S","sigma.T","b0","p","gamma","alpha")

# inits
inits <-list(N = (counts+1), b0 = rnorm(1,0,1),
             p = runif(1,0,1), 
             gamma = runif(1,0,1), 
             alpha = runif(1,0,50),
             eta = 0,
             rho = rnorm(1,0,1),
             S = rnorm(nrow(counts),0,1), 
             T = rnorm(ncol(counts),0,1),
             sigma.S = rtruncnorm(1,a=0,,mean=0,sd=1), sigma.T = rtruncnorm(1,a=0,,mean=0,sd=1),
             tau.S = pow(rtruncnorm(1,a=0,,mean=0,sd=1),-2), tau.T=pow(rtruncnorm(1,a=0,,mean=0,sd=1),-2))

# MCMC settings
nc <- 3
ni <- 8*10^4
nb <- ni*3/4
nt <- (ni-nb)/1000

# nimble model
Rmodel <- nimbleModel(code=model.ricker.nb, constants=constants, data=win.data, inits = inits)
Rmodel$initializeInfo()
# mcmc
conf <- configureMCMC(Rmodel, print = F)
conf$addMonitors(param.nb)
Rmcmc <- buildMCMC(conf) 
# Compile the model and MCMC algorithm
Cmodel <- compileNimble(Rmodel)
Cmcmc <- compileNimble(Rmcmc, project = Rmodel)

# Run the model
start.time <- Sys.time()
out <- runMCMC(mcmc = Cmcmc, niter = ni, nburnin = nb, nchains = nc, thin = nt,
               samplesAsCodaMCMC = T, inits = inits)
Sys.time()-start.time

MCMCvis::MCMCsummary(out, params = c("rho","eta","sigma.S","sigma.T","b0","p","gamma","alpha"))

mcmcplots::mcmcplot(out, parms = c("rho","eta","sigma.S","sigma.T","b0","p","gamma","alpha"))