#####################################################################################
#
# Example 4: Local level model with parameter learning
#
# Lopes and Tsay (2011)
# Particle Filters and Bayesian Inference in Financial Econometrics.
# Journal of Forecasting.
# 
#####################################################################################
#
# Comparison of MCMC and optimal auxiliary particle filter (state filtering)
# Comparison of MCMC and particle learning (state filtering and parameter learning)
#
# For t=1,...,n
#
#    y[t]|x[t]   ~ N(x[t],sig2)
#    x[t]|x[t-1] ~ N(alpha+beta*x[t-1],tau2)
#
# and
#
#    x[0]       ~ N(m0,C0)
#    sig2       ~ IG(a0,A0)
#    alpha|tau2 ~ N(b0[1],tau2*B0[1])
#    beta|tau2  ~ N(b0[2],tau2*B0[2])
#    tau2       ~ IG(nu0/2,nu0*tau20/2)
#
# Known quantities: (m0,C0,a0,A0,b0,B0,nu0,tau20)
#
######################################################################################
#
# Author : 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
#
#####################################################################################
rm(list=ls())

FFBS = function(y,alpha,beta,tau2,sig2,N){
  n    = length(y)
  b2   = beta^2
  B    = rep(0,n)
  H    = rep(0,n)
  mf   = rep(0,n)
  Cf   = rep(0,n)
  a    = rep(0,n+1)
  R    = rep(0,n+1)
  mb   = rep(0,n)
  Cb   = rep(0,n)
  a[1] = alpha+beta*m0
  R[1] = b2*C0+tau2
  for (t in 1:n){
    Q      = R[t]+sig2
    A      = R[t]/Q
    mf[t]  = (1-A)*a[t]+A*y[t]
    Cf[t]  = A*sig2
    a[t+1] = alpha+beta*mf[t]
    R[t+1] = b2*Cf[t]+tau2
    B[t]   = beta*Cf[t]/R[t+1]
    H[t]   = Cf[t]-B[t]^2*R[t+1]
  }
  mb[n] = mf[n]
  Cb[n] = Cf[n]
  for (t in (n-1):1){
    mb[t] = mf[t]+B[t]*(mb[t+1]-a[t+1])
    Cb[t] = Cf[t]-B[t]^2*(R[t+1]-Cb[t+1])
  }
  H     = sqrt(H)
  x     = matrix(0,n,N)
  x[n,] = rnorm(N,mf[n],Cf[n])
  for (t in (n-1):1)
    x[t,] = rnorm(N,mf[t]+B[t]*(x[t+1,]-a[t+1]),H[t])
  return(list(x=x,mf=mf,Cf=Cf,mb=mb,Cb=Cb))    
}

MCMC = function(y,alpha,beta,tau2,sig2,x,M){
  n  = length(y)
  xs = matrix(0,M,n)
  pars = matrix(0,M,4)
  for (i in 1:M){
    draw     = FFBS(y,alpha,beta,tau2,sig2,1)
    x        = as.numeric(draw$x)
    sig2     = 1/rgamma(1,n02+n/2,n0sig202+sum((y-x)^2)/2)
    y1       = x[2:n]
    X        = cbind(1,x[1:(n-1)])
    B1       = solve(diag(iB0)+t(X)%*%X)
    b1       = B1%*%(b0iB0+t(X)%*%y1)
    nu1tau21 = as.numeric(nu0tau20+t(y1-X%*%b1)%*%y1+t(b0-b1)%*%b0iB0)
    tau2     = 1/rgamma(1,nu02+(n-1)/2,nu1tau21/2)
    draw     = b1 + t(chol(tau2*B1))%*%rnorm(2)
    alpha    = draw[1]
    beta     = draw[2]
    pars[i,] = c(alpha,beta,tau2,sig2)
    xs[i,]   = x
  }
  return(list(pars=pars,xs=xs))  
}

PL = function(y,alphas,betas,tau2s,sig2s,xs){
  n      = length(y)
  N      = length(xs)
  draws  = array(0,c(n,N,5))
  s      = matrix(0,N,9)
  s[,1]  = 1.0/B0[1]
  s[,2]  = 0.0
  s[,3]  = 1.0/B0[2]
  s[,4]  = b0[1]/B0[1]
  s[,5]  = b0[2]/B0[2]
  s[,6]  = nu0
  s[,7]  = nu0*tau20 
  s[,8]  = n0
  s[,9]  = n0*sig20
  m      = s[,1]*s[,3]-s[,2]^2
  num1   = (s[,3]*s[,4]-s[,2]*s[,5])/m
  num2   = (s[,1]*s[,5]-s[,2]*s[,4])/m
  for (t in 1:n){
     # Resampling
     mus    = alphas+betas*xs
     stdevs = sqrt(sig2s+tau2s)
     weight = dnorm(y[t],mus,stdevs)
  	   k      = sample(1:N,size=N,replace=T,prob=weight)
  	   alphas = alphas[k]
  	   betas  = betas[k]
  	   tau2s  = tau2s[k]
  	   sig2s  = sig2s[k]
  	   mus    = mus[k]
  	   xs1    = xs[k]
  	   # Propagating
  	   vars   = 1/(1/sig2s+1/tau2s)
  	   sds    = sqrt(vars)
  	   means  = vars*(y[t]/sig2s + mus/tau2s)
  	   xs     = rnorm(N,means,sds)
  	  	 so     = s[k,]
    	 num1o  = num1[k]
  	   num2o  = num2[k]	  
  	   # Updating sufficient stats
  	   s[,1]  = so[,1] + 1
  	   s[,2]  = so[,2] + xs1
  	   s[,3]  = so[,3] + xs1^2
  	   s[,4]  = so[,4] + xs
  	   s[,5]  = so[,5] + xs*xs1
  	   s[,6]  = so[,6] + 1
     s[,8]  = so[,8] + 1
     s[,9]  = so[,9] + (y[t]-xs)^2
     # Sampling parameters
     sig2s  = 1/rgamma(N,s[8]/2,s[,9]/2)
  	   m      = s[,1]*s[,3]-s[,2]^2
  	   num1   = (s[,3]*s[,4]-s[,2]*s[,5])/m
  	   num2   = (s[,1]*s[,5]-s[,2]*s[,4])/m
  	   s[,7]  = so[,7] + (xs-num1-num2*xs1)*xs + (num1o-num1)*so[,4]+(num2o-num2)*so[,5]
     tau2s  = 1/rgamma(N,s[,6]/2,s[,7]/2)
     std    = sqrt(tau2s/m)
     norm   = cbind(rnorm(N,0,std),rnorm(N,0,std))
     alphas = num1 + sqrt(s[,3])*norm[,1]
     betas  = num2 - s[,2]/sqrt(s[,3])*norm[,1]+sqrt(s[,1]-s[,2]^2/s[,3])*norm[,2]  
     # Storing quantiles
     draws[t,,1] = alphas
     draws[t,,2] = betas
     draws[t,,3] = tau2s
     draws[t,,4] = sig2s
     draws[t,,5] = xs
  }
  return(draws)
}

PL.pure = function(y,alpha,beta,tau2,sig2,xs){
  n     = length(y)
  N     = length(xs)
  stdev = sqrt(sig2+tau2)
  vars  = 1/(1/sig2+1/tau2)
  sd    = sqrt(vars)
  xss   = matrix(0,n,N)
  for (t in 1:n){
    # Resampling
    mus = alpha+beta*xs
    w   = dnorm(y[t],mus,stdev)
    k   = sample(1:N,size=N,replace=T,prob=w)
    # Propagating
    means = vars*(y[t]/sig2+mus[k]/tau2)
    xs    = rnorm(N,means,sd)
    xss[t,] = xs
  }
  return(xss)
}

PL.bs = function(y,alphas,betas,tau2s,xs){
  N  = length(tau2s)
  xb = rep(0,n)
  xbs = matrix(0,n,N) 
  for (i in 1:N){
  	 print(i)
    alpha = alphas[i]
    beta = betas[i]
    tau = sqrt(tau2s[i])
    xb[n] = xs[n,i]
    for (t in (n-1):1){
      w = dnorm(xb[t+1],alpha+beta*xs[t,],tau)
      k = sample(1:N,size=1,prob=w)
      xb[t] = xs[t,k]
    }
    xbs[,i] = xb
  }
  return(xbs)
}

PL.pure.bs = function(y,alpha,beta,tau2,xs){
  n  = length(y)
  N  = ncol(xs)
  xb = rep(0,n)
  xbs = matrix(0,n,N) 
  for (i in 1:N){
  	 print(i)
    xb[n] = xs[n,i]
    for (t in (n-1):1){
      w = dnorm(xb[t+1],alpha+beta*xs[t,],tau)
      k = sample(1:N,size=1,prob=w)
      xb[t] = xs[t,k]
    }
    xbs[,i] = xb
  }
  return(xbs)
}
###################################################################################################################

# Simulated data
# --------------
set.seed(98765)
n     =  200
alpha =  0.0
beta  =  0.9
tau2  =  0.5
sig2  =  1.0
tau   = sqrt(tau2)
sig   = sqrt(sig2)
x     = rep(alpha/(1-beta),n+1)
true  = c(alpha,beta,tau2,sig2)
names = c("alpha","beta","tau2","sigma2")
for (t in 2:(n+1))
  x[t] = alpha+beta*x[t-1]+tau*rnorm(1)
x = x[2:(n+1)]
y = rnorm(n,x,sig)
xtrue = x


par(mfrow=c(1,1))
plot(y,ylim=range(x,y),xlab="Time",ylab="",main="",pch=16)
lines(x,col=2,lwd=2)

# Prior hyperparameters
# ---------------------
m0        = 0.0
C0        = 10
n0        = 10
sig20     = sig2
nu0       = 10
tau20     = tau2
b0        = c(alpha,beta)
B0        = c(1,1)/tau20
sC0       = sqrt(C0)
sB0       = sqrt(B0)
n02       = n0/2
n0sig202  = n0*sig20/2
nu02      = nu0/2
nu0tau20  = nu0*tau20
nu0tau202 = nu0*tau20/2
iB0       = 1/B0
b0iB0     = b0/B0

########################################
# PL and true density - fixed parameters
########################################
set.seed(246521)
qpure  = array(0,c(4,n,3))
N      = 2000
M      = 10000
mcmcf  = matrix(0,n,M)
xs     = rnorm(N,m0,sC0)
plpf   = PL.pure(y,alpha,beta,tau2,sig2,xs)
plpb   = PL.pure.bs(y,alpha,beta,tau2,plpf)
dlm    = FFBS(y,alpha,beta,tau2,sig2,M)
qtruef = cbind(dlm$mf+qnorm(0.025)*sqrt(dlm$Cf),dlm$mf,dlm$mf+qnorm(0.975)*sqrt(dlm$Cf))
qtrueb = cbind(dlm$mb+qnorm(0.025)*sqrt(dlm$Cb),dlm$mb,dlm$mb+qnorm(0.975)*sqrt(dlm$Cb))
qpure[1,,] = cbind(apply(plpf,1,quantile,0.025),apply(plpf,1,median),apply(plpf,1,quantile,0.975))
qpure[2,,] = cbind(apply(plpb,1,quantile,0.025),apply(plpb,1,median),apply(plpb,1,quantile,0.975))
qpure[3,,] = cbind(apply(dlm$x,1,quantile,0.025),apply(dlm$x,1,median),apply(dlm$x,1,quantile,0.975))
for (i in 10:n){
	print(i)
  mcmcf[i,]=FFBS(y[1:i],alpha,beta,tau2,sig2,M)$x[i,]
}
qpure[4,,] = cbind(apply(mcmcf,1,quantile,0.025),apply(mcmcf,1,median),apply(mcmcf,1,quantile,0.975))

pdf(file="filter-smoother-purefilter.pdf",width=10,height=10)
par(mfrow=c(2,1))
ts.plot(qtruef,xlab="Time",ylab="States",main="",ylim=range(qtruef,qtrueb,qpure),lwd=2)
for (i in 1:3) lines(qpure[1,,i],col=grey(0.5),lwd=2)
legend(10,4,legend=c("TRUE","OAPF"),col=c(1,grey(0.5)),lwd=c(2,2),lty=c(1,1),bty="n")
title("Forward filtering")
ts.plot(qtrueb,xlab="Time",ylab="States",main="",ylim=range(qtruef,qtrueb,qpure),lwd=2)
for (i in 1:3) lines(qpure[2,,i],col=grey(0.5),lwd=2)
legend(10,4,legend=c("TRUE","OAPF"),col=c(1,grey(0.5)),lwd=c(2,2),lty=c(1,1),bty="n")
title("Backward sampling")
dev.off()

#######################################
# PL and MCMC with parameter estimation
#######################################
# MCMC/FFBS
# ---------
set.seed(45498)
M0   = 1000
M    = M0+1000
time.mcmc = system.time({mcmc=MCMC(y,alpha,beta,tau2,sig2,xtrue,M)})
xsmcmc = t(mcmc$xs[(M0+1):M,])
qxmcmc = cbind(apply(xsmcmc,1,quantile,0.025),apply(xsmcmc,1,median),apply(xsmcmc,1,quantile,0.975))

times = NULL
ts    = seq(10,n,by=10)
for (t in ts){
  print(t)
  time.mcmc1=system.time({mcmc1=MCMC(y[1:t],alpha,beta,tau2,sig2,xtrue[1:t],M)})
  times = c(times,time.mcmc1[3])
}
par(mfrow=c(1,1))
plot(ts,times)
sum(cbind(1,1:n)%*%lm(times~ts)$coef)

# Particle filter and particle smoother
# -------------------------------------
set.seed(246521)
N      = 2000
xs     = rnorm(N,m0,sC0)
tau2s  = 1/rgamma(N,nu02,nu0tau202)
taus   = sqrt(tau2s)
alphas = rnorm(N,b0[1],taus*sB0[1])
betas  = rnorm(N,b0[2],taus*sB0[2])
sig2s  = 1/rgamma(N,n02,n0sig202)
pars   = cbind(alphas,betas,tau2s,sig2s)
par(mfrow=c(2,2))
for (i in 1:4){
  hist(pars[,i],prob=TRUE,xlab="",main=names[i])
  abline(v=true[i],col=2)
}
time.plf = system.time({pl.filter = PL(y,alphas,betas,tau2s,sig2s,xs)})
alphas   = pl.filter[n,,1]
betas    = pl.filter[n,,2]
tau2s    = pl.filter[n,,3]
sig2s    = pl.filter[n,,4]
xs       = pl.filter[,,5]
time.pls = system.time({pl.smoother = PL.bs(y,alphas,betas,tau2s,xs)})
qxpls    = cbind(apply(pl.smoother,1,quantile,0.025),apply(pl.smoother,1,median),apply(pl.smoother,1,quantile,0.975))

# Comparing MCMC/FFBS and PF
# --------------------------
pdf(file="particle-filter-smoother.pdf",width=10,height=10)
par(mfrow=c(2,1))
ts.plot(qtrueb,xlab="Time",ylab="States",main="",ylim=range(qtruef,qtrueb,qpure),lwd=2)
for (i in 1:3) lines(qpure[2,,i],col=grey(0.5),lwd=2)
legend(10,4,legend=c("TRUE","OAPF"),col=c(1,grey(0.5)),lwd=c(2,2),lty=c(1,1),bty="n")
title("(a)")
ind = 1:n
ts.plot(qxpls[ind,],xlab="Time",ylab="",main="",ylim=range(qxpls[ind,],qxmcmc[ind,]),lwd=2)
for (i in 1:3) lines(qxmcmc[ind,i],col=grey(0.5),lwd=2)
legend(10,4,legend=c("MCMC","PL"),col=c(1,grey(0.5)),lwd=c(2,2),lty=c(1,1),bty="n")
title("(b)")
dev.off()

pdf(file="parameters.pdf",width=7,height=12)
par(mfrow=c(4,2))
for (i in 1:4){
  serie1 = mcmc$pars[(M0+1):M,i]
  serie2 = pl.filter[n,,i]
  breaks = seq(min(serie1,serie2),max(serie1,serie2),length=20)
  h1 = hist(serie1,breaks=breaks,plot=FALSE)
  h2 = hist(serie2,breaks=breaks,plot=FALSE)
  U = max(h1$density,h2$density)
  hist(serie1,prob=TRUE,xlab="",main=names[i],ylab="MCMC",xlim=range(serie1,serie2),breaks=breaks,ylim=c(0,U))
  box()
  abline(v=true[i],col=2)
  hist(serie2,prob=TRUE,xlab="",main=names[i],ylab="PL",xlim=range(serie1,serie2),breaks=breaks)
  abline(v=true[i],col=2)
  box()
}
dev.off()