############################################################################################
# 
# SEQUENTIAL MONTE CARLO METHODS 
# 
# 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@ChicagoGSB.edu
# URL: http://faculty.chicagobooth.edu/hedibert.lopes/research/
#
############################################################################################
q025=function(x){quantile(x,0.025)}

q975=function(x){quantile(x,0.975)}

DLM = function(y,sig2,tau2,m0,C0){
  n  = length(y)
  m  = rep(0,n)
  C  = rep(0,n)
  mb = rep(0,n)
  Cb = rep(0,n)
  B  = rep(0,n-1)
  for (t in 1:n){
    if (t==1){
      a = m0
      R = C0 + tau2
    }else{
      a = m[t-1]
      R = C[t-1] + tau2
    }
    A = R/(R+sig2)
    m[t] = (1-A)*a + A*y[t]
    C[t] = A*sig2
    B[t-1] = C[t]/(C[t]+tau2)
  }
  mb[n] = m[n]
  Cb[n] = C[n]
  for (t in (n-1):1){
    mb[t] = (1-B[t])*m[t]+B[t]*mb[t+1]
    Cb[t] = (1-B[t])*C[t]+B[t]^2*Cb[t+1]
  }
  return(list(mf=m,Cf=C,mb=mb,Cb=Cb))
}

dlm.sim = function(n,sig,tau,x0){
  x    = rep(0,n)
  y    = rep(0,n)
  x[1] = rnorm(1,x0,tau)
  y[1] = rnorm(1,x[1],sig)
  for (t in 2:n){
    x[t] = rnorm(1,x[t-1],tau)
    y[t] = rnorm(1,x[t],sig)
  }
  return(list(x=x,y=y))
}

SISR = function(y,x,w,tau,sig,smooth){
  n  = length(y)
  N  = length(x)
  wf = matrix(0,n,N)
  xf = matrix(0,n,N)
  for (t in 1:n){
    x       = rnorm(N,x,tau)
    w       = dnorm(y[t],x,sig)
    x       = sample(x,size=N,replace=T,prob=w)
    xf[t,]  = x
    wf[t,]  = w
  }
  if(smooth){
    xb = matrix(0,n,N)
    for (j in 1:N){
    	xb[n,j] = xf[n,j]
    	for (t in (n-1):1){
    	  w = dnorm(xb[t+1,j],xf[t,],tau)
    	  xb[t,j] = sample(xf[t,],size=1,prob=w)
    	}
    }
    return(list(xf=xf,wf=wf,xb=xb))
  }else{
    return(list(xf=xf,wf=wf))
  }
}

APF = function(y,x,w,tau,sig,smooth){
  n  = length(y)
  N  = length(x)
  wf = matrix(0,n,N)
  xf = matrix(0,n,N)  
  for (t in 1:n){
    w   = dnorm(y[t],x,sig)
  	 k   = sample(1:N,size=N,replace=TRUE,prob=w)
    x1  = rnorm(N,x[k],tau)
    w   = dnorm(y[t],x1,sig)/dnorm(y[t],x[k],sig)
    x   = sample(x1,size=N,replace=T,prob=w)
    xf[t,]  = x
    wf[t,]  = w
  }
  if(smooth){
    xb = matrix(0,n,N)
    for (j in 1:N){
    	xb[n,j] = xf[n,j]
    	for (t in (n-1):1){
    	  w = dnorm(xb[t+1,j],xf[t,],tau)
    	  xb[t,j] = sample(xf[t,],size=1,prob=w)
    	}
    }
    return(list(xf=xf,wf=wf,xb=xb))
  }else{
    return(list(xf=xf,wf=wf))
  }
}


# Simulating 1st order DLM
# ------------------------
set.seed(123456)
n    = 50
tau2 = 0.5
sig2 = 1.0
x0   = 0.0
sig  = sqrt(sig2)
tau  = sqrt(tau2)
sim  = dlm.sim(n,sig,tau,x0)
x    = sim$x
y    = sim$y

# Prior hyperparameters
# ---------------------
m0 = 0
C0 = 100

# Closed form online estimation
# -----------------------------
dlm = DLM(y,sig2,tau2,m0,C0)
lxf = dlm$mf+qnorm(0.025)*sqrt(dlm$Cf)
mxf = dlm$mf
uxf = dlm$mf+qnorm(0.975)*sqrt(dlm$Cf)
lxb = dlm$mb+qnorm(0.025)*sqrt(dlm$Cb)
mxb = dlm$mb
uxb = dlm$mb+qnorm(0.975)*sqrt(dlm$Cb)

L = min(lx)
U = max(ux)
par(mfrow=c(1,1))
plot(mx,xlab="time",ylab="",ylim=c(L,U),lwd=2,axes=F,pch=16,type="l")
axis(2);axis(1,at=ind,label=ind);box()
lines(lx,lwd=2)
lines(ux,lwd=2)

# SISR
# ----
set.seed(12545)
N       = 1000
x       = rnorm(N,m0,sqrt(C0))
smooth  = TRUE
sisr    = SISR(y,x,w,tau,sig,smooth)
mxf.sisr = apply(sisr$xf,1,median)
lxf.sisr = apply(sisr$xf,1,q025)
uxf.sisr = apply(sisr$xf,1,q975)
mxb.sisr = apply(sisr$xb,1,median)
lxb.sisr = apply(sisr$xb,1,q025)
uxb.sisr = apply(sisr$xb,1,q975)

# Auxiliary particle filter (APF)
# -------------------------------
set.seed(12545)
N      = 1000
x      = rnorm(N,m0,sqrt(C0))
apf    = APF(y,x,w,tau,sig,smooth)
mxf.apf = apply(apf$xf,1,median)
lxf.apf = apply(apf$xf,1,q025)
uxf.apf = apply(apf$xf,1,q975)
mxb.apf = apply(apf$xb,1,median)
lxb.apf = apply(apf$xb,1,q025)
uxb.apf = apply(apf$xb,1,q975)

L = min(lx,lxf.sisr,lxf.apf,lxb.sisr,lxb.apf)
U = max(ux,uxf.sisr,lxf.apf,lxb.sisr,lxb.apf)
pdf(file="dlm-smoothing.pdf",width=10,height=7)
par(mfrow=c(2,2))
plot(mxf,xlab="time",ylab="",ylim=c(L,U),lwd=2,axes=F,pch=16,type="l")
title("SISR \n Forward filtering")
axis(2);axis(1,at=ind,label=ind);box()
lines(lxf,lwd=2)
lines(uxf,lwd=2)
lines(lxf.sisr,lwd=2,col=2)
lines(mxf.sisr,lwd=2,col=2)
lines(uxf.sisr,lwd=2,col=2)

plot(mxf,xlab="time",ylab="",ylim=c(L,U),lwd=2,axes=F,pch=16,type="l")
title("APF \n Forward filtering")
axis(2);axis(1,at=ind,label=ind);box()
lines(lxf,lwd=2)
lines(uxf,lwd=2)
lines(lxf.apf,lwd=2,col=2)
lines(mxf.apf,lwd=2,col=2)
lines(uxf.apf,lwd=2,col=2)

plot(mxb,xlab="time",ylab="",ylim=c(L,U),lwd=2,axes=F,pch=16,type="l")
title("SISR \n Backward sampling")
axis(2);axis(1,at=ind,label=ind);box()
lines(lxb,lwd=2)
lines(uxb,lwd=2)
lines(lxb.sisr,lwd=2,col=2)
lines(mxb.sisr,lwd=2,col=2)
lines(uxb.sisr,lwd=2,col=2)

plot(mxb,xlab="time",ylab="",ylim=c(L,U),lwd=2,axes=F,pch=16,type="l")
title("APF \n Backward sampling")
axis(2);axis(1,at=ind,label=ind);box()
lines(lxb,lwd=2)
lines(uxb,lwd=2)
lines(lxb.apf,lwd=2,col=2)
lines(mxb.apf,lwd=2,col=2)
lines(uxb.apf,lwd=2,col=2)
dev.off()