library(rstan)
library(dplyr)
library(psych)

rstan_options(auto_write = FALSE)
options(mc.cores = parallel::detectCores())
set.seed(123)

trawl_pooled_sp <- read.csv("TR_Catch_Sp.csv")
Zooplank = read.csv("ZooBio.csv")

uniqueyrs = unique(trawl_pooled_sp$Year)

uYear<-sort(unique(trawl_pooled_sp$Year))
orderyear<-sort(order(unique(trawl_pooled_sp$Year)))

usite<-sort(unique(trawl_pooled_sp$SITE))
ordersite<-sort(order(unique(trawl_pooled_sp$SITE)))

#assign year class and port order numbers
for (x in 1:length(uYear)){
  trawl_pooled_sp[which(trawl_pooled_sp[,1] == uYear[x]), 16]=orderyear[x]
}

for (x in 1:length(usite)){
  trawl_pooled_sp[which(trawl_pooled_sp[,2] == usite[x]), 17]=ordersite[x]
}

sdvals=trawl_pooled_sp %>%
  summarise(YEP=sd(log(YEP+1)),SPS=sd(log(SPS+1)),ALE=sd(log(ALE+1)),
            JOD=sd(log(JOD+1)),RAS=sd(log(RAS+1)),TRP=sd(log(TRP+1)),
            ROG=sd(log(ROG+1)),BLO=sd(log(BLO+1)),LOS=sd(log(LOS+1)),WHS=sd(log(WHS+1)))

trawl_pooled_sp %>%
  summarise(YEP=sd((YEP)),SPS=sd((SPS)),ALE=sd((ALE)),
            JOD=sd((JOD)),RAS=sd((RAS)),TRP=sd((TRP)),
            ROG=sd((ROG)),BLO=sd((BLO)),LOS=sd((LOS)),WHS=sd((WHS+1)))

trawl_pooled_sp %>%
  summarise(YEP=mean((YEP)),SPS=mean((SPS)),ALE=mean((ALE)),
            JOD=mean((JOD)),RAS=mean((RAS)),TRP=mean((TRP)),
            ROG=mean((ROG)),BLO=mean((BLO)),LOS=mean((LOS)),WHS=mean((WHS)))
## MCMC settings
ni = 6000
nt = 1
nb = 3000
nc = 3

functional = c(1,1,2,3,2,1,3,2,3,3)
ZoopScaled=scale(Zooplank$TotalBiomass)
datalist=list(N = nrow(trawl_pooled_sp),
              T = max(trawl_pooled_sp$yearn),
              sites = max(trawl_pooled_sp$siten),
              S = ncol(trawl_pooled_sp[,6:15]),
              N_func = max(functional),
              functional=functional,
              y = round(trawl_pooled_sp[,6:15],0),
              year=trawl_pooled_sp$yearn,
              site=trawl_pooled_sp$siten,
              meanalpha = log(as.numeric(trawl_pooled_sp[1,6:15])+1),
              monum=trawl_pooled_sp$monum,
              yearzoop = nrow(Zooplank),
              zoop = ZoopScaled[,1]
)


## Initial values
inits <- lapply(1:nc, function(i)
  list(  
    sigma_obs = sdvals +1,
    sigma_proc = runif(datalist$N_func,0.1,0.5),
    sigma_site = runif(datalist$S,0.5,1),#
    sitere_raw = matrix(rnorm(datalist$sites*datalist$S,0,1),nrow=datalist$sites),
    r_raw = matrix(runif(((datalist$T)-1)*datalist$S,-0.01,0),nrow=(datalist$T)-1),
    month_raw = matrix(rnorm(3*datalist$S,0,2),nrow=3),
    sigma_month = runif(datalist$S,0.5,1),
    init_log_alpha = log(trawl_pooled_sp[1,6:15]+1),
    sigmar = runif(1,0.1,1),
    meanr = matrix(runif(datalist$N_func*2,-0.1,0),nrow=2,ncol=datalist$N_func),
    rand = matrix(rep(0,datalist$N*datalist$S),nrow=(datalist$N)),
    ZooBeta = rnorm(datalist$N_func,0,0.5)
  )
)


m = stan_model(file = 'function_gps_model.stan')
out_2_4 = sampling(m,
                data = datalist, cores=nc,
                init = inits, #pars = params,
                chains = nc, iter = ni, warmup = nb, thin = nt,
                control = list(adapt_delta = 0.85, max_treedepth=20))

saveRDS(out_2_4, "Func_Out.rds")
q()