#============================================================================
#Load model inputs and create IDs
#============================================================================
PATH <- getwd()
load(paste0(PATH,"/","DataALL.RData"))
LakeID <- sort(unique(DataALL$Lake))

#Lakes where stocking is important to fish catch
ST.Lake<-  c("Balaton","Chaohu","Constance","Geneva","Kinneret","Michigan","Simcoe","Taihu") 

#ID for incomplete fishing effort and catch (Not included in the catch model)
EffID                     <- rep(1,1395)
EffID[is.na(DataALL$EFF)] <- 2

#Reorganize DataALL
DataALL <- DataALL[order(DataALL$Lake, EffID),]

#============================================================================
#Data for Bayesian networks model
#============================================================================
DataB <- list(
#Data and parameter
DFflat  = 4,    #DF for flat prior 

#Fishery variables [Y = Yield, easier to code]
logY = log(DataALL$CATCH), logST = log(DataALL$ST), logEFF = log(DataALL$EFF),

#Lake environment variables
WT = DataALL$WT, dWL = DataALL$dWL, logCHL = log(DataALL$CHL),

#Catchment climate and land use variables
AT = DataALL$AT, PRE = DataALL$PRE, PE = DataALL$PE, 
logPRE = log(DataALL$PRE), LUag = DataALL$LUag, logLUag = log(DataALL$LUag))

#============================================================================
# IDs and indices for MCMC simulations
#============================================================================
#Lakes with stocking
DataB$S0 <- c(1:31)[-which(unique(DataALL$Lake) %in% ST.Lake)]
#Lakes without stocking
DataB$S1 <- which(unique(DataALL$Lake) %in% ST.Lake)

#First data row of a lake
DataB$JS1 <- c(0, as.vector(cumsum(table(DataALL$Lake)))[1:30]) + 1
#Last data row of a lake for catch model
DataB$JE1 <- DataB$JS1 + c(matrix(table(DataALL$Lake,EffID),ncol = 2)[,1]) - 1 
#Last data row of a lake for lake environment model
DataB$JE2 <- DataB$JS1 + 44                                                            

#=============================================================================
#Lake by lake priors from linear models
# WT models
#=============================================================================
DataB$kT        <- 2
DataB$Beta.hatT <- matrix(rep(c(1:31),DataB$kT)  , ncol = DataB$kT)
DataB$sigma.hatT<- rep(-1,31)
DataB$VT        <- matrix(rep(c(1:31),DataB$kT^2), ncol = DataB$kT)
for (i in c(1:31)){
  SUM1 <- summary(lm(WT ~ AT, data = DataALL, na.action = "na.exclude", subset = Lake == LakeID[i]))
  DataB$Beta.hatT[i,]<- SUM1$coef[,1]
  DataB$sigma.hatT[i]<- SUM1$sigma
  DataB$VT[c(1:DataB$kT) + (i-1)*DataB$kT, ] <- SUM1$cov.unscaled}

#============================================================================
# dWL models
#============================================================================
DataB$kL        <- 4
DataB$Beta.hatL <- matrix(rep(c(1:31),DataB$kL)  , ncol = DataB$kL)
DataB$sigma.hatL<- rep(-1,31)
DataB$VL        <- matrix(rep(c(1:31),DataB$kL^2), ncol = DataB$kL)
for (i in c(1:31)){
  SUM1 <- summary(lm(dWL ~ PRE + PE + I(LUag * PE), data = DataALL, 
                     na.action = "na.exclude", subset = Lake == LakeID[i]))
  DataB$Beta.hatL[i,]<- SUM1$coef[,1]
  DataB$sigma.hatL[i]<- SUM1$sigma
  DataB$VL[c(1:DataB$kL) + (i-1)*DataB$kL, ] <- SUM1$cov.unscaled}

 
#============================================================================
# CHL models
#============================================================================
DataB$kC        <- 5
DataB$Beta.hatC <- matrix(rep(c(1:31),DataB$kC)  , ncol = DataB$kC)
DataB$sigma.hatC<- rep(-1,31)
DataB$VC        <- matrix(rep(c(1:31),DataB$kC^2), ncol = DataB$kC)
for (i in c(1:31)){
  SUM1 <- summary(lm(log(CHL) ~  log(PRE) + log(LUag) + dWL + WT, data = DataALL, 
                     na.action = "na.exclude", subset = Lake == LakeID[i]))
  DataB$Beta.hatC[i,]<- SUM1$coef[,1]
  DataB$sigma.hatC[i]<- SUM1$sigma
  DataB$VC[c(1:DataB$kC) + (i-1)*DataB$kC, ] <- SUM1$cov.unscaled}

#============================================================================
#CATCH model
#============================================================================
DataB$kYs        <- 6
DataB$Beta.hatYs <- matrix(rep(c(1:8),DataB$kYs)  , ncol = DataB$kYs)
DataB$sigma.hatYs<- rep(-1,8)
DataB$VYs        <- matrix(rep(c(1:8),DataB$kYs^2), ncol = DataB$kYs)

for (i in DataB$S1){
  SUM1 <- summary(lm(log(CATCH) ~ WT + dWL + log(CHL)  + log(ST) + log(EFF), 
                     data = DataALL, na.action = "na.exclude", subset = Lake == LakeID[i]))
  i <- which(DataB$S1 == i)
  DataB$Beta.hatYs[i,] <- SUM1$coef[,1]
  DataB$sigma.hatYs[i] <- SUM1$sigma
  DataB$VYs[c(1:DataB$kYs) + (i-1)*DataB$kYs, ] <- SUM1$cov.unscaled
}

DataB$kY        <- 5
DataB$Beta.hatY <- matrix(rep(c(1:23),DataB$kY)  , ncol = DataB$kY)
DataB$sigma.hatY<- rep(-1,23)
DataB$VY        <- matrix(rep(c(1:23),DataB$kY^2), ncol = DataB$kY)

for (i in DataB$S0){
  SUM1 <- summary(lm(log(CATCH) ~ WT + dWL + log(CHL) + log(EFF),
                     data = DataALL, na.action = "na.exclude", subset = Lake == LakeID[i]))
  i <- which(DataB$S0 == i)
  DataB$Beta.hatY[i,]<- SUM1$coef[,1]
  DataB$sigma.hatY[i]<- SUM1$sigma
  DataB$VY[c(1:DataB$kY) + (i-1)*DataB$kY, ] <- SUM1$cov.unscaled
}

#============================================================================
# Function for initial values
#============================================================================
BetaINI <- function(DataB = DataB, DataALL = DataALL){
library(MASS)
BetaV   <- function(Beta.hat, Vmat, sigma.hat, N = 2, Y = 31){
BetaOUT <- matrix(rep(c(1:Y),N), nrow = Y)
for (i in 1:Y){
  BetaOUT[i,] <- mvrnorm(1, Beta.hat[i,], Vmat[1:N, 1:N] * sigma.hat[i]^2)}
  return(BetaOUT)}
  
return(list(BetaT = BetaV(Beta.hat = DataB$Beta.hatT, Vmat = DataB$VT, sigma.hat = DataB$sigma.hatT, N = DataB$kT),
            BetaL = BetaV(Beta.hat = DataB$Beta.hatL, Vmat = DataB$VL, sigma.hat = DataB$sigma.hatL, N = DataB$kL),
            BetaC = BetaV(Beta.hat = DataB$Beta.hatC, Vmat = DataB$VC, sigma.hat = DataB$sigma.hatC, N = DataB$kC),
            BetaY = BetaV(Beta.hat = DataB$Beta.hatY, Vmat = DataB$VY, sigma.hat = DataB$sigma.hatY, N = DataB$kY, Y = 23),
            BetaYs= BetaV(Beta.hat = DataB$Beta.hatYs,Vmat = DataB$VYs,sigma.hat = DataB$sigma.hatYs,N = DataB$kYs,Y = 8)))}

#============================================================================
#Run JAGS
#============================================================================
MODELfile  <- paste0(PATH,"/","BNM.jag")
VAR1 <- c("BetaT","BetaL","BetaC","BetaY","BetaYs","Sig")

set.seed(888); library("rjags");library("runjags")
INI <- list(); for (i in c(1:3)){INI[[i]] <- BetaINI(DataB = DataB, DataALL = DataALL)}

Res0 <- run.jags(model = MODELfile, monitor = c(VAR1, "deviance"), data = DataB, inits = INI,
                 n.chains = 3, thin = 1, burnin = 10000, sample = 5 * 10^4, adapt = 10000, 
                 summarise = FALSE, method = "parallel")