#######################################################################
#
#  Example i. FIRST ORDER DLM 
#
#  AUXILIARY PARTICLE FILTER + PARAMETER LEARNING
#
#######################################################################
#
# Author : Hedibert Freitas Lopes
#          Graduate School of Business
#          University of Chicago
#          5807 South Woodlawn Avenue
#          Chicago, Illinois, 60637
#          Email : hlopes@ChicagoGSB.edu
#
#######################################################################
pri        = function(theta,theta0,sV0){dnorm(theta,theta0,sV0)}
likelihood = function(y,theta,sig){dnorm(y,theta,sig)}
rlike      = function(theta,sig){rnorm(1,theta,sig)}
post       = function(y,theta,theta0,V0){pri(theta,theta0,sV0)*likelihood(y,theta)}
quant025   = function(x){quantile(x,0.025)}
quant975   = function(x){quantile(x,0.975)}

# Data and prior hyperparameters 
theta0 = 0.0
V0     = 0.1
sV0    = sqrt(V0)
n0     = 5
S0     = 0.04

# Simulating the data
set.seed(13579)
n     = 200
tau2  = 0.01
sig2  = 0.02
sig2.true = sig2
tau   = sqrt(tau2)
sig   = sqrt(sig2)
y     = rep(0,n)
theta = rep(0,n)
theta[1] = 0.0
y[1] = rlike(theta[1],sig)
for (t in 2:n){
  theta[t] = rnorm(1,theta[t-1],tau)
  y[t] = rlike(theta[t],sig)
}

par(mfrow=c(1,1))
ts.plot(y,type="b")
lines(theta,col=2,type="b")

# Closed form solution
# --------------------
m     = rep(0,n)
C     = rep(0,n)
n1    = rep(0,n)
S1    = rep(0,n)
a     = theta0
R     = V0+tau2
f     = a
e     = y[1]-f
Q     = R + S0
A     = R/Q
n1[1] = n0 + 1
S1[1] = S0 + (S0/n1[1])*(e^2/Q-1)
m[1]  = a+A*e
C[1]  = A*S1[1]
for (t in 2:n){
  a = m[t-1]
  R = C[t-1]+tau2
  f = a
  e = y[t]-f
  Q = R + S1[t-1]
  A = R/Q
  n1[t] = n1[t-1] + 1
  S1[t] = S1[t-1] + (S1[t-1]/n1[t])*(e^2/Q-1)
  m[t] = a+A*e
  C[t] = A*S1[t]
}
Lm    = m + qt(0.025,n1)*sqrt(C)
Um    = m + qt(0.975,n1)*sqrt(C)
msig2 = (n1*S1/2)/(n1/2-1)
Ls1   = 1/qgamma(1-0.025,n1/2,n1*S1/2)
Us1   = 1/qgamma(1-0.975,n1/2,n1*S1/2)

# Auxiliarly particle filter with parameter learning
# --------------------------------------------------
N       = 10000
ws      = NULL
thetass = NULL
sig2s   = NULL
sig2    = 1/rgamma(N,n0/2,n0*S0/2)
thetas  = rnorm(N,theta0,sV0)
delta   = 0.75
h2       = 1-((3*delta-1)/(2*delta))^2
a       = sqrt(1-h2)
for (t in 1:n){
  lsig2   = log(sig2)
  mlsig2  = mean(lsig2)
  vlsig2  = var(lsig2)
  ms      = a*lsig2+(1-a)*mlsig2
  mus     = thetas 
  weight  = likelihood(y[t],mus,exp(ms/2))
  k       = sample(1:N,size=N,replace=T,prob=weight)
  sig2    = exp(rnorm(N,ms[k],sqrt(h2*vlsig2)))
  thetas  = rnorm(N,thetas[k],sqrt(tau2))
  weight  = likelihood(y[t],thetas,sqrt(sig2))/likelihood(y[t],mus[k],exp(ms[k]/2))
  ind     = sample(1:N,size=N,replace=T,prob=weight)
  thetas  = thetas[ind]
  sig2    = sig2[ind]
  ws      = rbind(ws,weight)
  sig2s   = rbind(sig2s,sig2)
  thetass = rbind(thetass,thetas)
}
mtheta  = apply(thetass,1,mean)
ltheta  = apply(thetass,1,quant025)
utheta  = apply(thetass,1,quant975)
msig2s  = apply(sig2s,1,mean)
lsig2s  = apply(sig2s,1,quant025)
usig2s  = apply(sig2s,1,quant975)

nn = NULL
for (i in 1:n)
  nn = c(nn,(100/N)*length(unique(thetass[i,])))
wss = apply(ws/matrix(apply(ws,1,sum),n,N),1,var)

# Graphical analysis
# ------------------
par(mfrow=c(2,2))
ts.plot(m,ylim=c(min(ltheta,Lm),max(utheta,Um)),ylab="",main=expression(theta[t]))
lines(m,col=2,lwd=1.25)
lines(Lm,col=2,lwd=1.25)
lines(Um,col=2,lwd=1.25)
lines(mtheta,col=4,lwd=1.25)
lines(ltheta,col=4,lwd=1.25)
lines(utheta,col=4,lwd=1.25)

ind=20:n
inds=seq(ind[1],ind[length(ind)],length=6)
plot(ind,msig2[ind],ylim=range(c(Ls1[ind],Us1[ind],lsig2s[ind],usig2s[ind])),ylab="",main=expression(sigma^2),type="l",xlab="",axes=F)
axis(2)
axis(1,at=inds,label=inds)
lines(ind,msig2[ind],col=2,lwd=1.25)
lines(ind,Ls1[ind],col=2,lwd=1.25)
lines(ind,Us1[ind],col=2,lwd=1.25)
lines(ind,msig2s[ind],col=4,lwd=1.25)
lines(ind,lsig2s[ind],col=4,lwd=1.25)
lines(ind,usig2s[ind],col=4,lwd=1.25)
abline(h=sig2.true,col=1,lwd=1.25)

ts.plot(nn,ylab="",main="Percentage of distinct particles",type="b")

ts.plot(wss,xlab="time",ylab="",main="Variance of weights",type="b")