library(terra)
library(sf)
library(exactextractr)
library(xts)
library(data.table)
library(ggplot2)
library(gridExtra)

# load data: precipitation, spi 1-month, spi 3-month, Kenya mask, population
prec <- rast("prec_kenya_monthly_1960_2023.nc")
spi1 <- rast("spi_kenya_monthly_1960_2023.nc", subds="spi1")
spi3 <- rast("spi_kenya_monthly_1960_2023.nc", subds="spi3")
mask <- vect("mask/ken_admbnda_adm0_iebc_20191031.shp")
population <- rast("population/ken_ppp_2020_0.05deg.nc")

# drought threshold - SPI below this value is considered to be a drought
threshold <- -1.0

# calculate time series of spatial average for visualization to see what the data looks like
data <- as.numeric(exact_extract(spi1, st_as_sf(mask), fun='mean'))
dates <- time(spi1)
spi1_ts <- xts(as.numeric(data), order.by=dates)
data <- as.numeric(exact_extract(spi3, st_as_sf(mask), fun='mean'))
dates <- time(spi3)
spi3_ts <- xts(as.numeric(data), order.by=dates)
# create a plot for SPI-1 and SPI-3
df <- data.frame("dates"=index(spi1_ts), "spi"=as.numeric(spi1_ts))
df <- as.data.table(df)
p1 <- ggplot(data=df, aes(x=dates,y=spi)) + 
	geom_col(data=df[spi>0], fill="blue") + 
	geom_col(data=df[spi<=0], fill="red") + 
	ggtitle(paste0("SPI-1 - Kenya")) +
	xlab(NULL) + ylab("SPI-1 [-]") +
	ylim(-3,3) +
	theme_bw()
df <- data.frame("dates"=index(spi3_ts), "spi"=as.numeric(spi3_ts))
df <- as.data.table(df)
p2 <- ggplot(data=df, aes(x=dates,y=spi)) + 
	geom_col(data=df[spi>0], fill="blue") + 
	geom_col(data=df[spi<=0], fill="red") + 
	ggtitle(paste0("SPI-3 - Kenya")) +
	xlab(NULL) + ylab("SPI-3 [-]") +
	ylim(-3,3) +
	theme_bw()
# plot the two plots together
#grid.arrange(p1, p2, ncol=1)

# identify drought periods
# for each month, if spi < drought threshold then set to 1, otherwise 0
spi1_threshold <- ifel(spi1 < threshold, 1, 0)
# calculate the population exposed to drought in each month
# start with a raster of population for when there is a drought
exposure <- ifel(spi1 < threshold, population, 0.0)
# sum the population exposed over the whole country
data <- as.numeric(exact_extract(exposure, st_as_sf(mask), fun='sum'))
# convert to a time series object
exposure_ts <- xts(as.numeric(data), order.by=dates)
# and plot to see how exposed population changes over time
plot(exposure_ts)