############################################################################################
# 
# FIRST ORDER DYNAMIC LINEAR MODEL
#
############################################################################################
#
# HEDIBERT FREITAS LOPES
# Associate Professor of Econometrics and Statistics
# The University of Chicago Booth School of Business
# 5807 South Woodlawn Avenue
# Chicago, Illinois, 60637
# Email : hlopes@ChicagoBooth.edu
# URL: http://faculty.chicagobooth.edu/hedibert.lopes/research/
#
############################################################################################
select = function(w){sample(1:M,size=1,prob=w)}
q025   = function(x){quantile(x,0.025)}
q975   = function(x){quantile(x,0.975)}

DLM = function(y,sig2,tau2,m0,C0,back){
  n  = length(y)
  m  = m0
  C  = C0
  l  = 0.0
  mf = rep(0,n)
  Cf = rep(0,n)
  for (t in 1:n){
    R = C + tau2
    Q = R + sig2
    l = l + dnorm(y[t],m,sqrt(Q),log=TRUE)
    A = R/Q
    m = (1-A)*m+A*y[t]
    C = A*sig2
    mf[t] = m
    Cf[t] = C
  }
  if(back){
    mb = rep(0,n)
    Cb = rep(0,n)
    mb[n] = mf[n]
    Cb[n] = Cf[n]
    for (t in (n-1):1){
 	   B = Cf[t]/(Cf[t]+tau2)
      mb[t] = (1-B)*mf[t]+B*mb[t+1]
      Cb[t] = (1-B)*Cf[t]+B*B*Cb[t+1]   
    }
    return(list(l=l,mf=mf,Cf=Cf,mb=mb,Cb=Cb))
  }else{
    return(list(l=l,mf=mf,Cf=Cf))
  }
} 

FFBS = function(y,sig2,tau2,m0,C0){
  n = length(y)
  m = rep(0,n)
  B = rep(0,n-1)
  H = rep(0,n-1)
  x = rep(0,n)
  for (t in 1:n){
    if (t==1){
      a = m0
      R = C0 + tau2
    }else{
      a      = m[t-1]
      R      = C + tau2
      B[t-1] = C/R
      H[t-1] = C*tau2/R
    }
    A = R/(R+sig2)
    m[t] = (1-A)*a + A*y[t]
    C = A*sig2
  }
  x[n] = rnorm(1,m[n],sqrt(C))
  if(n>1){
    for (t in (n-1):1)
      x[t] = (1-B[t])*m[t] + B[t]*x[t+1] + sqrt(H[t])*rnorm(1)
  }
  return(list(xb=x))
}

# Simulating 1st order DLM
# ------------------------
set.seed(123456)
tau2 = 0.5
sig2 = 1.0
n    = 100
x0   = 0.0
tau  = sqrt(tau2)
sig  = sqrt(sig2)
w    = rnorm(n,0,tau)
v    = rnorm(n,0,sig)
x    = rep(0,n)
y    = rep(0,n)
x[1] = x0 + w[1]
y[1] = x[1] + v[1]
for (t in 2:n){
  x[t] = x[t-1] + w[t]
  y[t] = x[t] + v[t]
}


pdf(file="dlm1.pdf",width=10,height=7)
plot(y,xlab="Time",type="b",ylab="",main="",pch=16,lwd=3,ylim=c(-5,10))
lines(x,col=2,lwd=3,type="b",pch=16)
legend(70,8,legend=c("y(t)","x(t)"),col=1:2,lwd=c(3,3),lty=c(1,1),bty="n")
dev.off()

# Prior information
m0 = 0.0
C0 = 10

dlm  = DLM(y,sig2,tau2,m0,C0,TRUE)

pdf(file="dlm2.pdf",width=10,height=7)
par(mfrow=c(1,1))
plot(dlm$mf,xlab="Time",ylab="",main="",ylim=c(-5,10),lwd=2,pch=16)
segments(1:n,dlm$mf+2*sqrt(dlm$Cf),1:n,dlm$mf-2*sqrt(dlm$Cf),lwd=2)
points(y,col=4,pch=16)
dev.off()

pdf(file="dlm3.pdf",width=10,height=7)
par(mfrow=c(1,1))
plot(dlm$mb,xlab="Time",ylab="",main="",ylim=c(-5,10),lwd=2,pch=16)
segments(1:n,dlm$mb+2*sqrt(dlm$Cb),1:n,dlm$mb-2*sqrt(dlm$Cb),lwd=2)
points(y,col=4,pch=16)
dev.off()


# Posterior of (sigma2,tau2)
nu0 = 5
s02 = 1
n0  = 5
t02 = 0.5
sig2s = seq(0.01,2,length=100)
tau2s = seq(0.01,2,length=100)
likes = matrix(0,100,100)
for (i in 1:100)
  for (j in 1:100)
    likes[i,j] = DLM(y,sig2s[i],tau2s[j],m0,C0,FALSE)$l+
                 dgamma(1/sig2s[i],nu0/2,nu0*s02/2,log=TRUE)-2*log(sig2s[i])+
                 dgamma(1/tau2s[j],n0/2,n0*t02/2,log=TRUE)-2*log(tau2s[j])

    
pdf(file="dlm4.pdf",width=10,height=7)
contour(sig2s,tau2s,exp(likes),xlab=expression(sigma^2),ylab=expression(tau^2),drawlabels=FALSE)
points(sig2,tau2,pch=16,cex=3)
dev.off()

#MCMC scheme
set.seed(202334)
M0 = 1000
M  = 2000
niter = M0+M
sig2d = sig2
tau2d = tau2
draws = matrix(0,niter,n+2)
for (iter in 1:niter){
  xd = FFBS(y,sig2d,tau2d,m0,C0)$xb
  par1 = (nu0+n)/2
  par2 = (nu0*s02+sum((y-xd)^2))/2
  sig2d = 1/rgamma(1,par1,par2)
  par1 = (n0+n-1)/2
  par2 = (n0*t02+sum((xd[2:n]-xd[1:(n-1)])^2))/2
  tau2d = 1/rgamma(1,par1,par2)
  draws[iter,] = c(sig2d,tau2d,xd)
}
ind = (M0+1):niter
draws = draws[ind,]


psig2 = apply(exp(likes),1,sum)
psig2 = psig2/sum(psig2)
hsig2 =sig2s[2]-sig2s[1]
psig2 = psig2/hsig2
ptau2 = apply(exp(likes),2,sum)
ptau2 = ptau2/sum(ptau2)
htau2 = tau2s[2]-tau2s[1]
ptau2 = ptau2/htau2

pdf(file="dlm5.pdf",width=12,height=7)
par(mfrow=c(1,3))
plot(draws[,1:2],xlab=expression(sigma^2),ylab=expression(tau^2),pch=16)
contour(sig2s,tau2s,exp(likes),drawlabels=FALSE,add=TRUE,col=2,lwd=2)
hist(draws[,1],xlab="",main=expression(sigma^2),ylim=range(psig2),
     breaks=seq(min(draws[,1]),max(draws[,1]),length=20),prob=TRUE)
lines(sig2s,psig2,col=2,lwd=2);box()
hist(draws[,2],xlab="",main=expression(tau^2),ylim=range(ptau2),
     breaks=seq(min(draws[,2]),max(draws[,2]),length=20),prob=TRUE)
lines(tau2s,ptau2,col=2,lwd=2);box()
dev.off()



sig2mode = mean(draws[,1])
tau2mode = mean(draws[,2])
dlm1  = DLM(y,sig2mode,tau2mode,m0,C0,TRUE)

xs = draws[,3:(n+2)]
qx = t(apply(xs,2,quantile,c(0.025,0.5,0.975)))

pdf(file="dlm6.pdf",width=15,height=7)
par(mfrow=c(1,1))
plot(qx[,2],xlab="Time",ylab="",main="",ylim=c(-5,10),pch=16)
segments(1:n,qx[,1],1:n,qx[,3])
points((1:n)-0.5,dlm1$mb,col=2,pch=16)
segments((1:n)-0.5,dlm1$mb+2*sqrt(dlm1$Cb),(1:n)-0.5,dlm1$mb-2*sqrt(dlm1$Cb),col=2)
dev.off()




# Alternative representation
Omega = matrix(0,n,n)
for (i in 1:n)
  for (j in i:n){
    Omega[i,j] = i
    Omega[j,i] = Omega[i,j]
  }    
I = diag(1,n)
II = matrix(1,n,n)
one = matrix(1,n,1)
VV = solve(II*C0+tau2*Omega)
V = solve(VV+I*sig2)
m = V%*%(VV%*%one*m0+y/sig2)

plot(m,type="l",ylim=c(-4,10))
lines(m-2*sqrt(diag(V)))
lines(m+2*sqrt(diag(V)))
lines(dlm$mb,col=2)
lines(dlm$mb+2*sqrt(dlm$Cb),col=2)
lines(dlm$mb-2*sqrt(dlm$Cb),col=2)