##Jake Ferguson 
#Formats environmental data, then  runs a gompertz population dynamic model on a number of beaver populations in JAG and save the output

library(R2jags)
library(dplyr)
library(RColorBrewer)
library(readxl)
library(coda)
holdout.prop	<- 0.0 #0.0 or 0.3 proportion of data that is held out of fitting to be used for prediction

##read in data
colony.dat 		<- as.data.frame(read_excel("Johnson-Bice_Ecol_App_forecasting_beaver_densities_data.xlsx", sheet=1, na="NA"))
envvars.dat 	<- as.data.frame(read_excel("Johnson-Bice_Ecol_App_forecasting_beaver_densities_data.xlsx", sheet=2, na="NA"))#[,-c(1,5)]
habvars.dat 	<- as.data.frame(read_excel("Johnson-Bice_Ecol_App_forecasting_beaver_densities_data.xlsx", sheet=3, na="NA"))

routename.vec   <- c("Blackduck", "Cass",  "Cass_crow", "C_st_louis", "Ely_finger", "Hay_kelliher", "Itasca", "Kabetogama", "Kanabec", "Kawishiwi", "Kooch_north", "Northome", "Red_lake", "Southern_pine", "West_vermillion") 

for(i in routename.vec) {
 temp <- colony.dat[which(colony.dat$rte.name == i),]$num.col
}
envvars.dat <- envvars.dat %>% filter(rte.name %in% routename.vec) %>% droplevels

habvars.dat		<- cbind(Route_name=habvars.dat$rte.name, data.frame(Quality=habvars.dat$relat_abund_good_forage, Quantity=habvars.dat$relative_forage_avail)) 

habvars.dat	 	<- habvars.dat[order(habvars.dat$Route_name),]

#extract appropriate covariates from the full table
curr.ind <- c(3, 32,32, 31,31, 16,16, 17,17, 18)
#set the appropriate lag values for each variable
lag.vals		<- c(1, 0,2, 2,3, 2,3, 2,3, 0, 0, 0)

envvars.dat		<- envvars.dat[, c(1,2,curr.ind)]
envvars.dat		<- cbind(envvars.dat, Quality=rep(habvars.dat$Quality, each=31))
envvars.dat 	<- cbind(envvars.dat[,1:2], scale(envvars.dat[,-(1:2)]))

year.min	<- min(colony.dat$year)
year.max 	<- max(colony.dat$year)
Dens.mat 	<- matrix(NA, year.max - year.min + 1, length(routename.vec))

dimnames(Dens.mat)[[1]] <- year.min:year.max
dimnames(Dens.mat)[[2]] <- routename.vec

Count.mat 		<- Observer.mat <- ObsType.mat <- km.mat <- Wolf.mat <- Dens.mat
init.index 		<- final.index <- vector('numeric', length(routename.vec))

EnvCov.array 	<- array(NA, c(year.max - year.min + 1, dim(envvars.dat)[2]-1, length(routename.vec))) 

dimnames(EnvCov.array)[[1]] <- year.min:year.max
dimnames(EnvCov.array)[[2]] <- c(dimnames(envvars.dat)[[2]][-c(1:2)], "Wolf")
dimnames(EnvCov.array)[[3]] <- routename.vec

holdout.obs <- vector('numeric', length(routename.vec))


##format input datasets for the jags model
for(i in 1:length(routename.vec)) {

	currCol 		<- colony.dat[which(colony.dat$rte.name == routename.vec[i]),]

	if(is.na(currCol$col.km[2])) {
		currCol 		  <- currCol[-(1:3),]
		curr.yearmin 	<- min(currCol$year)
		curr.yearmax 	<- max(currCol$year) 

	} else {
		curr.yearmin 	<- min(currCol$year)
		curr.yearmax 	<- max(currCol$year)
	}	

	init.index[i] 	<- curr.yearmin-year.min + 1
	final.index[i] 	<- curr.yearmax - year.min + 1
	curr.indices 	  <- init.index[i]:final.index[i]

	holdout.obs[i] 					<- round((final.index[i]-init.index[i]+1)*holdout.prop)

	Dens.mat[curr.indices, i]		<- currCol$col.km 
	Wolf.mat[curr.indices, i]		<- currCol$wolf_ubi
	Count.mat[curr.indices, i]	<- currCol$num.col
	km.mat[curr.indices, i]			<- currCol$rte.km 

	#format environmental covariates
	envTemp 	<- filter(envvars.dat, rte.name==routename.vec[i]) #filter for current route
	First.index	<- which(envTemp$year == curr.yearmin) - lag.vals
	Last.index 	<- which(envTemp$year == curr.yearmax) - lag.vals
	envTemp		<- envTemp[,-c(1,2)] #remove year and routename from covariate matrix

	temp <- matrix(NA, length(First.index[1]:Last.index[1]), dim(envTemp)[2])
	for(j in 1:dim(temp)[2]) {
		envVec	<- envTemp[First.index[j]:Last.index[j],j]
		if(sd(envVec) == 0){
			temp[,j] <- envVec
			next()
		} #one case is all zeros so can't use scale function
		temp[,j] <- envVec 
	}

	EnvCov.array[curr.indices,1:dim(temp)[2],i]	<-  temp #cast to matrix, scale covariates

}

#now add in wolf covariate
EnvCov.array[,dim(EnvCov.array)[2],] <- Wolf.mat

#scale covariates
for(i in 1:dim(EnvCov.array)[2]) {
	EnvCov.array[,2,] <- (EnvCov.array[,2,] - mean(EnvCov.array[,2,], na.rm=T))/sd(EnvCov.array[,2,], na.rm=T)
}

for(i in 1:dim(EnvCov.array)[2]) {
	dimnames(EnvCov.array)[[2]][i] <- paste(dimnames(EnvCov.array)[[2]][i], '-lag', lag.vals[i], sep='')
}

NumPops			<- length(routename.vec)
numEnvVars	<- dim(EnvCov.array)[2]
datTable 		<- list(lD=log(Dens.mat), NumPops=NumPops, init.index=init.index, final.index=final.index, EnvCov.array=EnvCov.array, NumEnvVars=numEnvVars, holdout.obs=holdout.obs)
overall.dens 	<- apply(Count.mat, 1, sum, na.rm=T)/apply(km.mat, 1, sum, na.rm=T)
l.dens 			<- log(overall.dens)


initTable 	<- function() {
	list(a=rnorm(1, -0.1, 0.1), b=rnorm(1, -0.1, 0.01))
}

par.est		<- c('a', 'sigmaInt', 'b', 'sigmaSlope', 'betaEnv', 'aRE', 'sigmaProc', 'Xpred')
jags.fit  	<- jags(model.file='Beavermodels.jags', n.iter=1e4, n.chains=4, data=datTable, inits=initTable, parameters.to.save=par.est)

#fit the model
jags.fit 		<- update(jags.fit, n.iter=1e5, n.thin=10) 

#calculate infomration criteria
cat("DIC", jags.fit$BUGSoutput$DIC, "\n")

load.module("dic")
m <- jags.model("Beavermodels.jags", data=datTable, inits=initTable, n.chains = 2)
update(m, 3000)
s <- jags.samples(m, c("deviance", "WAIC"), type="mean", n.iter=10000, thin=10)
s <- lapply(s, unclass)
cat('WAIC', sum(sapply(s, sum)), '\n')


jags.fit.mcmc 	<- as.mcmc(jags.fit)

jags.fit.summ 	<- summary(jags.fit.mcmc)[[1]]
names.vec 		  <- dimnames(jags.fit.summ)[[1]]

save.image(file=paste("Beaver_Holdout", holdout.prop, ".Rdata", sep=''))