####################################################################################################
#
# Modeling daily covid death with dynamic linear models
# 
# By Hedibert Lopes
# Professor of Statistics and Econometrics
# Head of the Center for Statistics, Data and Decision Sciences
# Insper Institute of Education and Research
# URL: www.hedibert.org
#
# March 9th 2021.
#
# Dynamic modeling package in R:
# https://www.unofficialgoogledatascience.com/2017/07/fitting-bayesian-structural-time-series.html
#
####################################################################################################

library("tidyverse")
library("readr")
library("forecast")
library("bsts") 


url = "https://github.com/CSSEGISandData/COVID-19/raw/master/csse_covid_19_data/csse_covid_19_time_series/"
url.deaths = paste(url,"time_series_covid19_deaths_global.csv",sep="")
deaths      = read_csv(url(url.deaths))
country    = "Brazil"
p = ncol(deaths)
y = diff(as.numeric(deaths[deaths[,2]==country,5:p]))
y = y[60:length(y)]
n = length(y)

# State space model
# ------------------------------------------------------------------------------------------------
ss   = AddLocalLinearTrend(list(),y)
ss   = AddSeasonal(ss, nseasons = 7,y)
ssm  = bsts(y,state.specification = ss,niter=10000,family="gaussian")

par(mfrow=c(1,1))
plot(ssm,"components",1)


h    = 55
pred = predict(ssm,horizon=h)
trendfit = apply(ssm$state.contributions[,1,],2,mean)
meanfit = apply(ssm$state.contributions[,1,]+ssm$state.contributions[,2,],2,mean)

pdf(file="trend-forecast.pdf",width=12,height=10)
par(mfrow=c(1,1))
plot(pred,xlab=paste("Last ",n," days",sep=""),ylab="Daily death",ylim=c(0,4000))
lines(trendfit,col=2,lwd=2)
lines(meanfit,col=4)
lines(c(n,n+1),c(meanfit[n],pred$mean[1]),col=4)
points(y,pch=16)
legend("topleft",legend=c("Observed counts","Trend","Fitted & forecasted counts","Credibility interval"),col=c(1,2,4,3),lty=1,lwd=2)
title("Simple dynamic linear model\n Forecast up to April 30th 2021 (another 100K deaths)")
abline(h=1000,lty=2)
abline(h=2000,lty=2)
abline(h=3000,lty=2)
abline(v=n)
axis(1,at=n,lab="March 6th,2021")
dev.off()



qs.t = t(apply(ssm$state.contributions[,1,],2,quantile,c(0.25,0.5,0.75)))
qs.s = t(apply(ssm$state.contributions[,2,],2,quantile,c(0.25,0.5,0.75)))

n = length(y)
mmovel = rep(0,n)
mmovel[1:6] = y[1:6]
for (t in 7:n)
  mmovel[t] = mean(y[t:(t-6)])
ratio1 = 100*(mmovel[14:n]/mmovel[1:(n-13)]-1)

ratios = matrix(0,10000,n)
for (j in 1:10000)
  ratios[j,14:n] = 100*(ssm$state.contributions[j,1,14:n]/ssm$state.contributions[j,1,1:(n-13)]-1)
qr = t(apply(ratios,2,quantile,c(0.25,0.5,0.75)))


par(mfrow=c(1,1))
ts.plot(qs.t)
lines(mmovel,col=2)
abline(h=1000,col=4,lwd=3)

pdf(file="trend-growthrate.pdf",width=12,height=10)
par(mfrow=c(2,1),mai=c(0.9,0.9,0.3,0.3))
plot(qs.t[(n-130):n,2],type="l",ylim=range(qs.t[(n-130):n,]),axes=FALSE,
     xlab="Most recent day is Saturday, March 6th 2021",ylab="Death trend")
lines(qs.t[(n-130):n,1],lty=2)
lines(qs.t[(n-130):n,3],lty=2)
axis(2);box();axis(1,at=1:131,lab=130:0)
abline(h=1000,col=4,lwd=3)

plot(qr[(n-130):n,2],type="l",axes=FALSE,ylim=range(qr[(n-130):n,]),
     xlab="Most recent day is Saturday, March 6th 2021",ylab="14-day growth rate")
lines(qr[(n-130):n,1],lty=2)
lines(qr[(n-130):n,3],lty=2)
axis(2);box();axis(1,at=1:131,lab=130:0)
abline(h=15,col=4,lwd=3)
abline(h=-15,col=4,lwd=3)
dev.off()