source("functions.R")
library("zoo")
library("lubridate")
library("Matrix")
library("pracma")
library("ggpubr")
library("gridGraphics")

## Load data
N <- read.table("data/NDVI.txt")
N <- as.numeric(N$V1)
fill_value <- -999
# climate conditions in order: precipitation plus snowmelt (mm/day), 
# air temperature (degree C), vapor pressure deficit (kPa),
# net shortwave radiation (W/m2)
CLM <- read.table('data/CLM.txt') 
# daily averages of climate conditions in the same order
AvgCLM <- read.table('data/AvgCLM.txt')

## Compute climate anomaly within each interval of two NDVI observations
date0 <- as.Date("1999-06-30") # the data of first NDVI obervation

deltaT <- difftime(date0, as.Date("1999-06-14"))
anCLM <- data.frame(V1 = 0, V2 = 0, V3 = 0, V4 = 0)
for (i in 0:length(CLM[,1])){
  st <- date0 + deltaT*i
  st <- strftime(st, format = "%j")
  et <- date0 + deltaT*(i+1)
  et <- strftime(et, format = "%j")
  
  if (et < st){
    window <- c((as.numeric(st)+1):365, 1:et)
  }else{
    window <- c((as.numeric(st)+1):et)
  }
  window[window==365] <- 0  # leap year
  anCLM <- rbind(anCLM, CLM[i+1,] - colMeans(AvgCLM[window,], na.rm = TRUE))
}

anCLM <- anCLM[2:(nrow(anCLM)-1),]

## Center NDVI time series
N[N==fill_value]<-NA
Y <- N[2:length(N)] - mean(N, na.rm = T)

## Use two seasonal harmonic components
rseas <- c(1,2) 

## Include lag-1 centered NDVI and precipitation in the regression module
X <- cbind( (N - mean(N, na.rm = T))[1:(length(N)-1)],     #lag-1 NDVI
       anCLM[1:(nrow(anCLM)-1),1] )                   # precipitation 

## Defining model inputs
delta <- 0.98

M <- list()
M$Y <- Y
M$X <- X
M$rseas <- rseas
dd <- rep(delta, 4) 
M$delta <- dd

ntrend <- 2; nregn <- ncol(X); pseas <- length(rseas); nseas <- pseas*2
dimen <- ntrend+nregn+nseas
m <- rep(0, dimen)
C <- matrix(0, nrow = dimen, ncol = dimen)
diag(C) <- 1
S <- 0.2^2
nu <- ntrend+nregn+pseas
pr <- list(m, C, S, nu)

M$prior <- pr

## Running model
FF <- forwardFilteringM(M)

## Model likelihood
slik <- FF[[4]]
## Extract estimates of the coefficient corresponding to lag-1 NDVI
sm <- FF[[1]]
ind <- 3                      # index of autocorrelation
sm.m <- as.numeric(sm[ind,])  # mean of autocorrelation

sC <- FF[[2]]                 
sC.v <- sC[[ind]][ind,]       # variance of autocorrelation

snu <- FF[[3]]                # degrees of freedom

## Plotting the results
tsN <- ts(N, start = c(1999,06,30), frequency=23)
quantile1 <- 0.95
quantile2 <- 0.80

lown <- Index_low(N, date0, quantile2) # timesteps when NDVI was abnormally low
# Now find periods when it was low for at least half of the time (3 steps) 
# in three months (6 steps).
lown_cont <- c()
for (i in 1:length(lown)){
  tile <- c()
  for (j in lown){
    if(j<=lown[i] & j>=lown[i]-5)
      tile <- c(tile, j)
  }
  if(length(tile)>=3){
    lown_cont <- c(lown_cont, lown[i])
  }
}

plot(tsN, type="p", axes=F, pch=16, xlab="", ylab="", ylim=c(0.35,0.8))
rect(time(tsN)[180], 0, time(tsN)[249], 1, 
     col = "#FF003322", border = "white")
rect(time(tsN)[295], 0, time(tsN)[length(tsN)], 1, 
     col = "#FF003322", border = "white")
points(tsN, pch=16)
lines(tsN, lwd=1)
axis(1, cex.axis=1.2)
axis(2, cex.axis=1.2)
title(xlab="Time", ylab="NDVI", cex.lab=1.4)
points(tsN[lown_cont]~time(tsN)[lown_cont], col="red", pch=16)
legend("bottomleft", legend = c("NDVI", "ALN"), col = c("black", "red"),
       pch=16, lty = c(1,0), cex=1.3, ncol=2)

##### Second plot #####
steps <- seq(date0, as.Date("2015-12-03"), by = 16)
diebackdate <- steps[lown_cont[1]]
warmup <- 47 # two years (46 time steps) warmup period

bd <- sm.m + sC.v * qt(quantile1, snu) # 0.95 boundary
bd2 <- sm.m + sC.v * qt(quantile2, snu) # 0.8 boundary

mbd <- median(bd[(warmup+1):length(bd)], na.rm = T) # above this line, considered a EWS
ews <- which(sm.m>mbd) # finding time steps of EWS
ews <- ews[ews>warmup] # only outside of the warmup period

# Selecting periods where EWS was above threshold for more than 3 months
ews_continuous <- c()
window <- floor(90/16) # three months

for (i in 1:length(ews)){
  tile <- c()
  for (j in ews){
    if(j<=ews[i] & j>=ews[i]-window)
      tile <- c(tile, j)
  }
  if(length(tile)>window-1){
    ews_continuous <- c(ews_continuous, ews[i])
  }
}

tmp <- steps[ews_continuous]
ewsdate <- tmp[tmp>as.Date("2012-07-15")][1] # arbitrary?
mortdate <- as.Date("2015-07-15") # from external dataset

# Plotting
library(ggplot2)

smEWS<-sm.m
smEWS[-ews_continuous]<-NA
df <- data.frame(Autocorrelation = sm.m, Time = steps, bd = bd, bd2 = bd2, 
                 ews = smEWS)
## Axes and background
gg1<- ggplot(df, aes(x = Time, y = Autocorrelation))+
  annotate("rect", ymin=-1.5, ymax=1.6, xmin=steps[173], xmax=steps[240], 
           fill="red", alpha=0.1)+
  annotate("rect", xmin=steps[287], xmax=steps[length(steps)], ymin=-1.5, ymax=1.6,
           alpha=0.1, fill="red")+
  theme_classic()+
  xlim(steps[1], steps[length(steps)])+
  scale_y_continuous(expand = c(0,0), limits = c(-1.5, 1.6))+
  theme(axis.text = element_text(size=12), axis.title = element_text(size=16))+
  theme(axis.title.y = element_text(margin = margin(t = 0, r = 10, b = 0, l = 0)))
## Lines and points
gg2 <- gg1 + geom_ribbon(aes(ymin=2*Autocorrelation-bd, ymax=bd, fill="lightgray"))+
  geom_ribbon(aes(ymin=2*Autocorrelation-bd2, ymax=bd2, fill="darkgray"))

gg3 <- gg2 + geom_line(aes(color="darkblue"))+
  geom_point(aes(y=ews, color="red"), shape=17)+
  geom_segment(aes(x=Time[warmup], xend=Time[length(Time)], y=mbd, yend=mbd),
               linetype="dashed")
## Legend
gg4 <- gg3 + coord_cartesian(ylim=c(-0.5,0.6)) + 
  scale_fill_manual("", values = c("darkgray", "lightgray"), labels=c("80% range","95% range"))+
  scale_colour_manual("", values = c("darkblue","red"), labels=c("Mean","EWS"),
                      guide = guide_legend(override.aes = list(
                        linetype = c("solid", "blank"),shape = c(NA,17))))+
  guides(fill=guide_legend(nrow=2))+
  theme(legend.position = c(0.5,0.9), legend.text=element_text(size=12), 
        legend.box = "horizontal", legend.title = element_blank(),
        legend.background = element_blank(),legend.spacing.y = unit(0, "mm"),
        legend.box.background = element_rect(colour="black"))
gg4