#################################################################################################
#
#  NON-LINEAR NORMAL DYNAMIC MODEL 
#
#  COMPARING SISR AND APF
#
#################################################################################################
#
# Author : Hedibert Freitas Lopes
#          The University of Chicago Booth School of Business
#          5807 South Woodlawn Avenue
#          Chicago, Illinois, 60637
#          Email : hlopes@Chicagobooth.edu
#
#################################################################################################
rm(list=ls())

q025 = function(x){
  quantile(x,0.025)
}

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

g = function(arg){
  c(arg[1],arg[1]/(1+arg[1]^2),cos(1.2*arg[2]))
}

fullxt = function(t,yt,xtm1,xt,xtp1,theta,sig,tau){
  zm1  = g(c(xtm1,t-1))
  z    = g(c(xt,t))
  full = dnorm(xt,sum(zm1*theta),tau,log=TRUE)+dnorm(xtp1,sum(z*theta),tau,log=TRUE)+dnorm(yt,xt^2/20,sig,log=TRUE)
  return(full)
}

fullxn = function(t,yt,xtm1,xt,theta,sig,tau){
  zm1  = g(c(xtm1,t-1))
  z    = g(c(xt,t))
  full = dnorm(xt,sum(zm1*theta),tau,log=TRUE)+dnorm(yt,xt^2/20,sig,log=TRUE)
  return(full)
}

draw.x = function(y,x,x0,theta,sig,tau,v){
  n   = length(y)
  x1  = rnorm(1,x[1],v)
  num = fullxt(1,y[1],x0,x1,x[2],theta,sig,tau)
  den = fullxt(1,y[1],x0,x[1],x[2],theta,sig,tau)
  if (log(runif(1))<(num-den)){
    x[1]=x1
  }
  xn  = rnorm(1,x[n],v)
  num = fullxn(n,y[n],x[n-1],xn  ,theta,sig,tau)
  den = fullxn(n,y[n],x[n-1],x[n],theta,sig,tau)
  if (log(runif(1))<(num-den)){
    x[n]=xn
  } 
  for (t in 2:(n-1)){
    xt  = rnorm(1,x[t],v)
    num = fullxt(t,y[t],x[t-1],xt  ,x[t+1],theta,sig,tau)
    den = fullxt(t,y[t],x[t-1],x[t],x[t+1],theta,sig,tau)
    if (log(runif(1))<(num-den)){
      x[t]=xt
    }
  }
  return(x)
}

fixedpar = function(y,X,b,A,v,lam){
  n     = length(y)
  k     = ncol(X)
  par1  = (v+n)/2
  var   = solve(crossprod(X,X)+A)
  mean  = matrix(var%*%(crossprod(X,y)+crossprod(A,b)),k,1)
  par2  = v*lam + sum((y-crossprod(t(X),mean))^2)
  par2  = (par2 + crossprod(t(crossprod(mean-b,A)),mean-b))/2
  sig2  = 1/rgamma(1,par1,par2)
  var   = var*sig2
  mean  = mean + crossprod(t(chol(var)),rnorm(k))
  return(c(mean,sig2))
}
#################################################################################################################

# Simulating the data
set.seed(54678)
n         = 100
sig2      = 1
tau2      = 10
sig       = sqrt(sig2)
tau       = sqrt(tau2)
theta     = c(0.5,25,8)
ptr       = c(theta,tau2,sig2)
y         = rep(0,n)
x         = rep(0,n)
x0        = 0.1
Z         = g(c(x0,0))
x[1]      = rnorm(1,sum(Z*theta),tau)
y[1]      = rnorm(1,x[1]^2/20,sig)
for (t in 2:n){
  Z    = g(c(x[t-1],t-1))
  x[t] = rnorm(1,sum(Z*theta),tau)
  y[t] = rnorm(1,x[t]^2/20,sig)
}
xtr = x

#pdf(file="nonlinear.pdf",width=10,height=7)
par(mfrow=c(2,2))
ts.plot(y,xlab="time",ylab="",main=expression(y[t]))
ts.plot(x,xlab="time",ylab="",main=expression(x[t]))
plot(x,y,xlab=expression(x[t]),ylab=expression(y[t]))
plot(x[1:(n-1)],x[2:n],xlab=expression(x[t-1]),ylab=expression(x[t]))
#dev.off()

# Prior hyperparameters
# ---------------------
m0 = 0.0
C0 = 10.0

# Brute force sequential MCMC
# ---------------------------
set.seed(122545)
v     = 0.1
M0    = 1000
M     = 5000
niter = M0+M
xss   = NULL
ns    = seq(10,n,by=10)
for (t in ns){
  x     = xtr[1:t]
  xs    = NULL
  for (iter in 1:niter){	
    x  = draw.x(y[1:t],x,x0,theta,sig,tau,v)
    xs = c(xs,x[t])
  }
  xss = rbind(xss,xs)
}
x.smcmc   = xss
mx.smcmc = apply(x.smcmc[,(M0+1):niter],1,median)
Lx.smcmc = apply(x.smcmc[,(M0+1):niter],1,q025)
Ux.smcmc = apply(x.smcmc[,(M0+1):niter],1,q975)

# SISR
# ----
set.seed(12545)
N   = 5000
x   = rnorm(N,m0,sqrt(C0))
xs  = NULL
for (t in 1:n){
  X  = cbind(x,x/(1+x^2),cos(1.2*(t-1)))
  mx = X%*%theta
  x1 = rnorm(N,mx,tau)
  w  = dnorm(y[t],x1^2/20,sig)
  x   = sample(x1,size=N,replace=T,prob=w)
  xs  = rbind(xs,x)
}
x.sisr  = xs
mx.sisr = apply(xs,1,median)
Lx.sisr = apply(xs,1,q025)
Ux.sisr = apply(xs,1,q975)

# APF
# ---
set.seed(12545)
N   = 5000
x   = rnorm(N,m0,sqrt(C0))
xs  = NULL
for (t in 1:n){
  X   = cbind(x,x/(1+x^2),cos(1.2*(t-1)))
  mx  = X%*%theta
  w   = dnorm(y[t],mx^2/20,sig)
  ind = sample(1:N,size=N,replace=T,prob=w)
  x1  = rnorm(N,mx[ind],tau)
  w1  = dnorm(y[t],x1^2/20,sig)/w[ind]
  x   = sample(x1,size=N,replace=T,prob=w1)
  xs  = rbind(xs,x)
}
x.apf  = xs
mx.apf = apply(xs,1,median)
Lx.apf = apply(xs,1,q025)
Ux.apf = apply(xs,1,q975)

# Comparison 
# ----------
ind = ns
L = min(Lx.smcmc,Lx.sisr[ind],Lx.apf[ind])
U = max(Ux.smcmc,Ux.sisr[ind],Ux.apf[ind])

pdf(file="nonlinear-smcmc-sisr-apf.pdf",height=5,width=10)
par(mfrow=c(1,2))
plot(ind,mx.smcmc,xlab="Time",ylab="",ylim=range(L,U),type="l",lwd=3,col=grey(0.6))
lines(ind,Lx.smcmc,lwd=3,col=grey(0.6))
lines(ind,Ux.smcmc,lwd=3,col=grey(0.6))
lines(ind,Lx.sisr[ind],lwd=2,col=2)
lines(ind,mx.sisr[ind],lwd=2,col=2)
lines(ind,Ux.sisr[ind],lwd=2,col=2)
title("TRUE vs. SISR")

plot(ind,mx.smcmc,xlab="Time",ylab="",ylim=range(L,U),type="l",lwd=3,col=grey(0.6))
lines(ind,Lx.smcmc,lwd=3,col=grey(0.6))
lines(ind,Ux.smcmc,lwd=3,col=grey(0.6))
lines(ind,Lx.apf[ind],lwd=2,col=2)
lines(ind,mx.apf[ind],lwd=2,col=2)
lines(ind,Ux.apf[ind],lwd=2,col=2)
title("TRUE vs. APF")
dev.off()