###Part I: Preparing the data
rm(list=ls())
getwd()
#Load all the necessary packages (need to download and install them first)

library(rethinking)
library(rstan)
library(nlme)
library(bbmle)
library(lme4)     # for fitting GLMMs
library(lattice)  # for the xyplot function
library(MuMIn)    # for the r.squaredGLMM function

##Set your working directory manually or using the following commands## 
#getwd()
#setwd("pathway")
#setwd("C:/Users/pquintan/Documents/Classes/Methods in Ecology II/Demos/")

## read the Hypericum data from file and prepare data 

orig_data <- read.table("hypericum_data_94_07.txt", header=T)
dist_data <- read.table("Bald_mat_dis.txt", header=T)
coord_data <- read.table("Bald_coordinates.txt", header=T)
coord_data <- as.data.frame(coord_data)
dist_data <- as.matrix(dist_data/1000,14,14)
dist <- as.data.frame(dist_data)
round(dist_data,2)
site <- unique(orig_data$bald)
ran_sites <-sample(site,6)
ran_sites <- ran_sites[order(ran_sites)]

dis_pop_ran <- array(0,c(6,6))
coo_pop_ran <- array(0,c(6,3))
coo_pop_ran[,1] <-ran_sites 
for(i in 1:6){
  coo_pop_ran[i,2] <-as.numeric(coord_data[which(ran_sites[i]==site),2])
  coo_pop_ran[i,3] <-as.numeric(coord_data[which(ran_sites[i]==site),3])
  for(j in 1:6){
    dis_pop_ran[i,j] <- dist[which(ran_sites[i]==site), 
                             which(ran_sites[j]==site)]
  }  
}
colnames(dis_pop_ran) <- ran_sites
rownames(dis_pop_ran) <- ran_sites
dis_pop_ran 

dt <- subset(orig_data, !is.na(ht_init) & !is.na(st_init) & rp_init > 0 & 
               year==1995 & bald %in% ran_sites )  ### EXPLAIN THE CHANGE HERE  
yr <- unique(dt$year)

dt$lgh <- log(dt$ht_init)
dt$lfr <- log(dt$rp_init)

table(dt$bald,dt$fire_year)
dt$tsf <- dt$year-dt$fire_year
table(dt$bald,dt$tsf)

dt$lgh_s <- scale(dt$lgh)
dt$tsf_s <- scale(dt$tsf)

######################  RANDOM MODELS ##############################################

colSums(is.na(dt))
dt<-dt[,colSums(is.na(dt))==0]
dt$stage<-as.factor(dt$stage)
dt$fate<-as.factor(dt$fate)
## model with no random effects

m_no <- ulam(
  alist(
    lfr ~ dnorm(mu,sigma),
    mu <- a + b*lgh_s + c*tsf_s + cc*tsf_s*lgh_s,
    a ~ dnorm(0,50),
    b ~ dnorm(0,1),
    c ~ dnorm(0,1),
    cc ~ dnorm(0,1),
    sigma ~ dunif(0,1)
  ),
  data = dt,chains =3,
  iter=6000,warmup=2000, log_lik = TRUE,
)

precis(m_no,digits=3)
plot(m_no,pars=c("a","b","c","cc","sigma"))
# pairs(m_no,pars=c("a","b","c","cc","sigma"))
pairs(m_no,horInd=c(1:5),verInd=c(1:5))
summary(lm(lfr ~ lgh_s * tsf_s, data=dt)) 

postn <- extract.samples(m_no)

par(mfrow=c(1,6))
for(i in names(postn)){
  dfpostn <- data.frame(postn[i])
  plot(as.numeric(rownames(dfpostn)), dfpostn[,1], type = "l",
       xlab = "Index",
       ylab = i)
}


# Assign populations indices
dt$pj <- 1
for (i in 2: length(ran_sites)){
  dt$pj[dt$bald==ran_sites[i]] <-i}

# model with random intercepts
m_rint <- ulam(
  alist(
    lfr ~ dnorm(mu,sigma),
    mu <- a + p[pj]+ b*lgh_s + c*tsf_s + cc*tsf_s*lgh_s,
    a ~ dnorm(0,50),
    b ~ dnorm(0,1),
    c ~ dnorm(0,1),
    cc ~ dnorm(0,1),
    p[pj] ~ dnorm(0,sigmap),
    sigmap ~ dcauchy(0,1),
    sigma ~ dcauchy(0,1)
  ),
  data = dt,chains =3, log_lik = TRUE,
)

precis(m_rint,depth=2,digits=3,pars=c("a","b","c","cc","sigma"))
precis(m_rint,depth=2,digits=3,pars=c("p","sigmap"))
par(mfrow=c(1,1))
plot(m_rint,pars=c("a","b","c","cc","sigma"))
plot(m_rint, depth=2)
# pairs(m_rint,pars=c("a","b","c","cc","sigma"))
pairs(m_no,horInd=c(1:5),verInd=c(1:5))
summary(lme(lfr~lgh_s*tsf_s,random=~1|bald,data=dt,method="REML"))

postp <- extract.samples(m_rint)

par(mfrow=c(2,4))
for(i in names(postp)){
  dfpostp <- data.frame(postp[i])
  plot(as.numeric(rownames(dfpostp)), dfpostp[,1], type = "l",
       xlab = "Index",
       ylab = i)
}

# model with random intercepts by space

m_rspac <- ulam(
  alist(
    lfr ~ dnorm(mu,sigma),
    mu <- a + p[pj]+ b*lgh_s+ c*tsf_s  + cc*tsf_s*lgh_s,
    transpars> vector[6]: p <<- L_SIGMA * z,
    vector[6]: z ~ dnorm(0 , 1),
    transpars> matrix[6,6]: L_SIGMA <<- cholesky_decompose(SIGMA),
    transpars> matrix[6,6]: SIGMA <- cov_GPL2(Dmat, etasq, rhosq, sigmap),
    # vector[6]:p ~ multi_normal(0 , SIGMA),
    # matrix[6,6]:SIGMA <- cov_GPL2( Dmat , etasq , rhosq , sigmap),
  #  p[pj] ~ GPL2(Dmat,etasq,rhosq,sigmap),
    a ~ dnorm(0,50),
    b ~ dnorm(0,10),
    c ~ dnorm(0,1),
    cc ~ dnorm(0,1),
    etasq ~ dcauchy(0,1),
    rhosq ~ dcauchy(0,1),
    sigmap ~ exponential(1),
    sigma ~ dcauchy(0,1)
  ),
  data = list(lfr = dt$lfr,
              lgh_s = dt$lgh_s,
              tsf_s = dt$tsf_s,
              pj = dt$pj,
              Dmat = dis_pop_ran),
  warmup = 2000, iter = 20000, log_lik = TRUE,
  chains =3)

precis(m_rspac,depth=2,digits=3)
par(mfrow=c(1,1))
plot(m_rspac,pars=c("a","b","c","cc","sigma"))
pairs(m_rspac, horInd=c(7:20),verInd=c(7:20))
pairs(m_rspac, horInd=c(12),verInd=c(12))

postr <- extract.samples(m_rspac)
par(mfrow=c(2,7))
for(i in names(postr)){
  dfpostr <- data.frame(postr[i])
  plot(as.numeric(rownames(dfpostr)), dfpostr[,1], type = "l",
       xlab = "Index",
       ylab = i)
}

compare(m_no,m_rint,m_rspac)

###### Evaluating the covariance among populations.
x <- seq(1,10,0.5)
post <- extract.samples(m_rspac)
par(mfrow=c(1,1),mar=c(4,4,2,2),mgp=c(1.6,0.6,0))
curve (median(post$etasq)*exp(-median(post$rhosq)*x^2),xlab="distance",
       ylab="covariance",ylim=c(0,1.5),lwd=2,xlim=c(0,10))
for (i in 1:100){
  curve (post$etasq[i]*exp(-post$rhosq[i]*x^2),add=TRUE,
         col=col.alpha("black",0.2))
}

unique(dt$tsf_s)
unique(dt$lgh_s) 
max(dt$lgh_s)
par(mfrow=c(1,1))
table(dt$tsf,dt$tsf_s)

plot(coo_pop_ran[,2],coo_pop_ran[,3],type="n", ylab= "x",xlab="y",
     xlim=c(-.1,0.4))
for (k in 1:6){
  d.sur <- data.frame(lgh_s=0,
                      pj=rep(k),
                      tsf_s=1.3775 ,
                      Dmat = 0
  )
  
  S.ensamble <- ensemble(m_rspac,data=d.sur)
  pred_S <- apply(S.ensamble$link,2,mean)
  #print(pred_S)
  points(coo_pop_ran[k,2],coo_pop_ran[k,3],
         cex=pred_S,pch=16)
}





#################################################ps
m.text <- c("no random","random Intercepts","spatail covariance")
par(mfrow=c (2,2))
plot(density(post$a),ylim=c(0,10), main ="a")
lines(density (postp$a) ,
      col="red")
lines(density (postn$a) ,
      col="blue")
legend(0.5,9.5, legend=m.text, pch= c(1,1,1), lty=1, col=c("black","red","blue"),cex=0.8)


plot(density(post$b),ylim=c(0,10), main ="b")
lines(density (postp$b) ,
      col="red")
lines(density (postn$b) ,
      col="blue")

plot(density(post$c),ylim=c(0,10), main ="c")
lines(density (postp$c) ,
      col="red")
lines(density (postn$c) ,
      col="blue")

plot(density(post$cc),ylim=c(0,10), main ="cc")
lines(density (postp$cc) ,
      col="red")
lines(density (postn$cc) ,
      col="blue")



#####################################
m_rspac_noint <- ulam(
  alist(
    lfr ~ dnorm(mu,sigma),
    mu <- a + p[pj]+ b*lgh_s+ c*tsf_s,
    transpars> vector[6]: p <<- L_SIGMA * z,
    vector[6]: z ~ dnorm(0 , 1),
    transpars> matrix[6,6]: L_SIGMA <<- cholesky_decompose(SIGMA),
    transpars> matrix[6,6]: SIGMA <- cov_GPL2(Dmat, etasq, rhosq, sigmap),
  #  p[pj] ~ GPL2(Dmat,etasq,rhosq,sigmap),
    a ~ dnorm(0,50),
    b ~ dnorm(0,1),
    c ~ dnorm(0,1),
    etasq ~ dcauchy(0,1),
    rhosq ~ dcauchy(0,1),
    sigmap ~ exponential(1),
    sigma ~ dcauchy(0,1)
  ),
  data = list(lfr = dt$lfr,
              lgh_s = dt$lgh_s,
              tsf_s = dt$tsf_s,
              pj = dt$pj,
              Dmat = dis_pop_ran),
              warmup = 6000, iter = 2e4, log_lik = TRUE, chains =3)

precis(m_rspac_noint,depth=2,digits=3)
par(mfrow=c(1,1))
plot(m_rspac_noint,depth=2)
pairs(m_rspac, horInd=c(7:20),verInd=c(7:20))
pairs(m_rspac, horInd=c(12),verInd=c(12))

post3 <- extract.samples(m_rspac_noint)
par(mfrow=c(2,6))
for(i in names(post3)){
  dfpost3 <- data.frame(post3[i])
  plot(as.numeric(rownames(dfpost3)), dfpost3[,1], type = "l",
       xlab = "Index",
       ylab = i)
}

compare(m_rint,m_rspac,m_rspac_noint)

#####################################
m_rspac_noint_nofire <- ulam(
  alist(
    lfr ~ dnorm(mu,sigma),
    mu <- a + p[pj]+ b*lgh_s,
    transpars> vector[6]: p <<- L_SIGMA * z,
    vector[6]: z ~ dnorm(0 , 1),
    transpars> matrix[6,6]: L_SIGMA <<- cholesky_decompose(SIGMA),
    transpars> matrix[6,6]: SIGMA <- cov_GPL2(Dmat, etasq, rhosq, sigmap),
   # p[pj] ~ GPL2(Dmat,etasq,rhosq,sigmap),
    a ~ dnorm(0,50),
    b ~ dnorm(0,1),
    etasq ~ dcauchy(0,1),
    rhosq ~ dcauchy(0,1),
    sigmap ~ exponential(1),
    sigma ~ dcauchy(0,1)
  ),
  data = list(lfr = dt$lfr,
              lgh_s = dt$lgh_s,
              pj = dt$pj,
              Dmat = dis_pop_ran),
              warmup = 6000, iter = 2e4, log_lik = TRUE,
              chains =3)

precis(m_rspac_noint_nofire,depth=2,digits=3)
par(mfrow=c(1,1))
plot(m_rspac_noint_nofire,)
pairs(m_rspac_noint_nofire, horInd=c(7:20),verInd=c(7:20))

compare(m_rint,m_rspac,m_rspac_noint,m_rspac_noint_nofire)

post4 <- extract.samples(m_rspac_noint_nofire)
par(mfrow=c(2,6))
for(i in names(post4)){
  dfpost4 <- data.frame(post4[i])
  plot(as.numeric(rownames(dfpost4)), dfpost4[,1], type = "l",
       xlab = "Index",
       ylab = i)
}

par(mfrow=c (1,1))
plot(density(post$b),ylim=c(0,15),xlim=c(0.8,1.4), main ="b")
lines(density (post3$b) ,
      col="red")
lines(density (post4$b) ,
      col="blue")

plot(density(post$rhosq),ylim=c(0,1),xlim=c(-20,20), main ="rho")
lines(density (post3$rhosq) ,
      col="red")
lines(density (post4$rhosq) ,
      col="blue")

### Corretations amomg populations

K <- array(0,c(6,6))
for(i in 1:6){
  for(j in 1:6){
    K[i,j] <- median(post$etasq)*exp(-median(post$rhosq)*dis_pop_ran[i,j]^2)
    diag(K) <- median(post$etasq)+0.01
  }
}
Rho <- round(cov2cor(K),2)
colnames(Rho) <- ran_sites
rownames(Rho) <- ran_sites
Rho