#################################################################################################
#
#  Example iii. NON-LINEAR NORMAL DYNAMIC MODEL 
#
#  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
#
#################################################################################################
fX         = function(arg){c(arg[1],arg[1]/(1+arg[1]^2),cos(1.2*arg[2]))}
pri        = function(theta,theta0,sV0){dnorm(theta,theta0,sV0)}
likelihood = function(y,theta,sig){dnorm(y,theta^2/20,sig)}
rlike      = function(theta,sig){rnorm(1,theta^2/20,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)}

# Simulating the data
set.seed(12345)
n         = 200
sig2      = 10
tau2      = 1
sig       = sqrt(sig2)
tau       = sqrt(tau2)
beta      = c(0.5,25,8)
pars.true = c(beta,sig2,tau2)
y         = rep(0,n)
theta     = rep(0,n)
theta0    = 0.1
X         = c(theta0,theta0/(1+theta0^2),0)
theta[1]  = rnorm(1,sum(X*beta),tau)
y[1]      = rlike(theta[1],sig)
for (t in 2:n){
  X        = fX(c(theta[t-1],t-1))
  theta[t] = rnorm(1,sum(X*beta),tau)
  y[t]     = rlike(theta[t],sig)
}

par(mfrow=c(2,2))
ts.plot(y,xlab="time",ylab="",main=expression(y[t]))
ts.plot(theta,xlab="time",ylab="",main=expression(theta[t]))
plot(theta,y,xlab=expression(theta[t]),ylab=expression(y[t]))
plot(theta[1:(n-1)],theta[2:n],xlab=expression(theta[t-1]),ylab=expression(theta[t]))

# Prior hyperparameters
# ---------------------
V0  = 0.1
sV0 = sqrt(V0)
a0  = 3
A0  = 20
b0  = 3
B0  = 2
c0  = c(0.5,25,8)
C0  = c(0.1,16,2)
m0  = 0
V0  = 5

# Auxiliary particle filter
# -------------------------
set.seed(8642)
N       = 1000
delta   = 0.75
h2      = 1-((3*delta-1)/(2*delta))^2
a       = sqrt(1-h2)
pars    = cbind(rnorm(N,c0[1],sqrt(C0[1])),rnorm(N,c0[2],sqrt(C0[2])),rnorm(N,c0[3],sqrt(C0[3])),
               log(1/rgamma(N,a0,A0)),log(1/rgamma(N,b0,B0)))
thetas  = rnorm(N,m0,sqrt(V0))
ESS     = NULL
thetass = NULL
parss   = array(0,c(N,5,n))
ws      = NULL
par(mfrow=c(1,1))
for (t in 1:n){
  mpar   = apply(pars,2,mean)
  vpar   = var(pars)
  ms     = a*pars+(1-a)*matrix(mpar,N,5,byrow=T)
  X      = t(apply(cbind(thetas,t-1),1,fX))
  mus    = apply(X*pars[,1:3],1,sum)
  weight = likelihood(y[t],mus,exp(ms[,4]/2))
  k      = sample(1:N,size=N,replace=T,prob=weight)
  ms1    = ms[k,] + matrix(rnorm(5*N),N,5)%*%chol(h2*vpar)
  X      = t(apply(cbind(thetas[k],t-1),1,fX))
  thetat = rnorm(N,apply(X*ms1[,1:3],1,sum),exp(ms1[,5]/2))
  w      = likelihood(y[t],thetat,exp(ms1[,4]/2))/likelihood(y[t],mus[k],exp(ms[k,4]/2))
  w      = w/sum(w)
  ind    = sample(1:N,size=N,replace=T,prob=w)
  thetas = thetat[ind]
  pars   = ms1[ind,]
  thetass     = rbind(thetass,thetas)
  parss[,,t]  = pars 
  ws          = rbind(ws,w)
  cv2         = var(w)/(mean(w)^2)
  ESS         = c(ESS,N/(1+cv2))
  ts.plot(ESS,xlim=c(1,n),ylim=c(0,N))
}

mtheta = apply(thetass,1,mean)
ltheta = apply(thetass,1,quant025)
utheta = apply(thetass,1,quant975)
mpars  = matrix(0,n,5)
lpars  = matrix(0,n,5)
upars  = matrix(0,n,5)

for (i in 1:3){
  mpars[,i] = apply(parss[,i,],2,mean)
  lpars[,i] = apply(parss[,i,],2,quant025)
  upars[,i] = apply(parss[,i,],2,quant975)
}
for (i in 4:5){
  mpars[,i] = apply(exp(parss[,i,]),2,mean)
  lpars[,i] = apply(exp(parss[,i,]),2,quant025)
  upars[,i] = apply(exp(parss[,i,]),2,quant975)
}

par(mfrow=c(2,3))
plot(theta,mtheta,pch=16,col=2,xlab=expression(theta),ylab=expression(E(theta)),xlim=c(-20,20),ylim=c(-20,20))
abline(0,1,col=4)
for (i in 1:n){
  segments(theta[i],ltheta[i],theta[i],utheta[i],col=3,lwd=1)
  points(theta[i],mtheta[i],col=2,pch=16)
}

names = c("alpha","beta","gamma","sigma2","tau2")
for (i in 1:5){
  ts.plot(lpars[,i],ylim=range(lpars[,i],upars[,i]),ylab="",main=names[i])
  lines(lpars[,i],col=1)
  lines(mpars[,i],col=1)
  lines(upars[,i],col=1)
  abline(h=pars.true[i],col=2)
}