###################################################################################
#
# Using Steve Scoot's bsts R package for Bayesian Structural Time Series modeling
# http://people.ischool.berkeley.edu/~hal/Papers/2013/pred-present-with-bsts.pdf
#
###################################################################################
install.packages("bsts")
library(bsts)

pdf(file="ss-via-bstspackage.pdf",width=10,height=5)

##########################################################################################
# PART I
# Local level model
# Linear growth model
# Dynamic regression model  
##########################################################################################

# SIMULATING SOME DATA
#########################
set.seed(1235)
n = 200
n1 = 75
n2 = 150
alpha1 = 1
alpha2 = 3
alpha3 = 0
beta1 = 1
beta2 = 1
beta3 = 0
sig = 0.5
x   = rnorm(n,2,1)
y1  = rnorm(n1,alpha1+beta1*x[1:n1],sig)
y2  = rnorm(n2-n1,alpha2+beta2*x[(n1+1):n2],sig)
y3  = rnorm(n-n2,alpha3+beta3*x[(n2+1):n],sig)
y = c(y1,y2,y3)
ind = c(rep(1,n1),rep(2,n2-n1),rep(4,n-n2))

par(mfrow=c(1,3))
plot(x,y)
ts.plot(y)
ts.plot(x)

# PART I.A) LOCAL LEVEL MODEL
#############################
M0   = 1000
M    = 9000
MM0  = M0+M
comp = AddLocalLevel(list(),y)
LLmodel = bsts(y, state.specification = comp, niter = MM0)

names(LLmodel)

dim(LLmodel$state.contributions)

qs = t(apply(LLmodel$state.contributions[(M0+1):MM0,1,],2,
     quantile,c(0.05,0.5,0.95)))

par(mfrow=c(1,2))
hist(LLmodel$sigma.obs[(M0+1):MM0],prob=TRUE,xlab="",main="")
title("Variance of obs. equation")
hist(LLmodel$sigma.level[(M0+1):MM0],prob=TRUE,xlab="",main="")
title("Variance of system equation")

par(mfrow=c(1,1))
ts.plot(qs,ylim=range(qs,y),xlab="Time",ylab="")
points(y,pch=16,col=2,cex=0.5)
legend("topleft",legend=c("y(t)","90%CI"),col=2:1,lty=1)


# PART I.B) LINEAR GROWTH MODEL
###############################
comp = AddLocalLinearTrend(list(), y)
LGmodel = bsts(y, state.specification = comp, niter = MM0)
qs1 = t(apply(LGmodel$state.contributions[(M0+1):MM0,1,],2,
      quantile,c(0.05,0.5,0.95)))

par(mfrow=c(2,3))
hist(LLmodel$sigma.obs[(M0+1):MM0],prob=TRUE,xlab="",main="")
title("Variance of obs. equation\nLocal level model")
hist(LGmodel$sigma.obs[(M0+1):MM0],prob=TRUE,xlab="",main="")
title("Variance of obs. equation\nLinear growth model")
hist(LLmodel$sigma.level[(M0+1):MM0],prob=TRUE,xlab="",main="")
title("Variance of system equation\Local level model")
hist(LGmodel$sigma.trend.level[(M0+1):MM0],prob=TRUE,xlab="",main="")
title("Variance of system equation 1\nLinear growth model")
hist(LGmodel$sigma.trend.slope[(M0+1):MM0],prob=TRUE,xlab="",main="")
title("Variance of system equation 2\nLinear growth model")


par(mfrow=c(1,2))
plot(density(LLmodel$sigma.obs[(M0+1):MM0]),xlab="",main="")
lines(density(LGmodel$sigma.obs[(M0+1):MM0]),col=2)
title("Variance of obs. equation")
legend("topright",legend=c("Local level","Linear Growth"),col=1:2,lty=1)
plot(density(LLmodel$sigma.level[(M0+1):MM0]),xlab="",main="",xlim=c(0,0.7))
lines(density(LGmodel$sigma.trend.level[(M0+1):MM0]),col=2)
title("Variance of system equation 1")


par(mfrow=c(1,1))
ts.plot(qs,ylim=range(qs,qs1,y),xlab="Time",ylab="",col=2)
lines(qs1[,1],col=3)
lines(qs1[,2],col=3)
lines(qs1[,3],col=3)
points(y,pch=16,cex=0.5)
legend("topleft",legend=c("y(t)","90%CI (LL)","90%CI (LG)"),col=1:3,lty=1)


# PART I.C) DYNAMIC REGRESSION MODEL
####################################
ss = list()
ss = AddLocalLevel(ss, y)
ss = AddDynamicRegression(ss, y ~ x)
model = bsts(y, state.specification = ss, niter = MM0)
qs1 = t(apply(model$state.contributions[(M0+1):MM0,1,],2,
      quantile,c(0.05,0.5,0.95)))
qs2 = t(apply(model$state.contributions[(M0+1):MM0,2,],2,
      quantile,c(0.05,0.5,0.95)))


par(mfrow=c(1,3))
hist(model$sigma.obs[(M0+1):MM0],prob=TRUE,xlab="",main="")
title("Variance of obs. equation")
hist(model$sigma.obs[(M0+1):MM0],prob=TRUE,xlab="",main="")
title("Variance of system equation\nDynamic intercept")
hist(model$V1.sigma[(M0+1):MM0],prob=TRUE,xlab="",main="")
title("Variance of system equation\nDynamic slope")

par(mfrow=c(1,2))
ts.plot(qs1,xlab="Time",ylab="",col=c(grey(0.8),1,grey(0.8)))
segments(1,alpha1,n1,alpha1,col=2)
segments(n1+1,alpha2,n2,alpha2,col=2)
segments(n2+1,alpha3,n,alpha3,col=2)
title("Time-varying intercept")
ts.plot(qs2,xlab="Time",ylab="",col=c(grey(0.8),1,grey(0.8)))
segments(1,beta1,n1,beta1,col=2)
segments(n1+1,beta2,n2,beta2,col=2)
segments(n2+1,beta3,n,beta3,col=2)
title("Time-varying slope")

##########################################################################################
# PART II
# Linear growth model with seasonal components
# Dynamic regression model  
##########################################################################################
data(AirPassengers)
y = log(AirPassengers)

par(mfrow=c(1,1))
plot(y,main="Airline data")

# PART II.A) Linear trend and seasonal dummies
############################################################
ss = AddLocalLinearTrend(list(), y)
ss = AddSeasonal(ss, y, nseasons = 12)
model = bsts(y, state.specification = ss, niter = MM0)

sigmas = cbind(model$sigma.obs,model$sigma.trend.level,
model$sigma.trend.slope,model$sigma.seasonal.12)

names = c("Obs. equation","State equation\nlevel",
"State equation\nslope","State equation\nseasonal")
par(mfrow=c(2,2))
for (i in 1:4)
  hist(sigmas[(M0+1):MM0,i],xlab="",main=names[i],prob=TRUE)


qs1 = t(apply(model$state.contributions[(M0+1):MM0,1,],2,quantile,c(0.05,0.5,0.95)))
qs2 = t(apply(model$state.contributions[(M0+1):MM0,2,],2,quantile,c(0.05,0.5,0.95)))

par(mfrow=c(1,2))
ts.plot(qs1,ylim=range(qs1,y),main="Trend component")
points(as.vector(y),pch=16,cex=0.5,col=2)
ts.plot(qs2,main="Sesonal component")

############################################################
# Model with linear trend and trigonometric seasonality
############################################################
ss <- AddLocalLinearTrend(list(), y)
ss <- AddTrig(ss, y, period = 12, frequencies = 1)
model1 = bsts(y, state.specification = ss, niter = 10000)

sigmas1 = cbind(
model1$sigma.obs,
model1$sigma.trend.level,
model1$sigma.trend.slope,
model1$trig.coefficient.sd[,1],
model1$trig.coefficient.sd[,2])

names = c("Obs. equation","State equation\nlevel",
"State equation\nslope","State equation\nseasonal 1","State equation\nseasonal 2")
par(mfrow=c(2,3))
for (i in 1:5)
  hist(sigmas1[(M0+1):MM0,i],xlab="",main=names[i],prob=TRUE)


qs1 = t(apply(model1$state.contributions[(M0+1):MM0,1,],2,quantile,c(0.05,0.5,0.95)))
qs2 = t(apply(model1$state.contributions[(M0+1):MM0,2,],2,quantile,c(0.05,0.5,0.95)))

par(mfrow=c(1,2))
ts.plot(qs1,ylim=range(qs1,y),main="Trend component")
points(as.vector(y),pch=16,cex=0.5,col=2)
ts.plot(qs2,main="Sesonal component")

qs11 = t(apply(model$state.contributions[(M0+1):MM0,1,]+model$state.contributions[(M0+1):MM0,2,],2,quantile,c(0.05,0.5,0.95)))
qs22 = t(apply(model1$state.contributions[(M0+1):MM0,1,]+model1$state.contributions[(M0+1):MM0,2,],2,quantile,c(0.05,0.5,0.95)))


par(mfrow=c(1,1))
ts.plot(qs11[,c(1,3)],ylim=range(qs11,qs22),main="",col=2)
points(as.vector(y),pch=16,cex=0.5)
lines(qs22[,1],col=4)
lines(qs22[,3],col=4)
legend("topleft",legend=c("y","Free form seasonality","Trigonometric seasonality"),col=c(1,2,4),lty=1)

dev.off()