#################################################################################################
#
#  NON-LINEAR NORMAL DYNAMIC MODEL 
#
#################################################################################################
#
# 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())
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){
  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))
}

quant025   = function(x){quantile(x,0.025)}
quant975   = function(x){quantile(x,0.975)}

# 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
n0      = 6
sig20   = 2/3
theta0  = c(0.5,25,8)
V0      = diag(c(0.25,10,4))/tau2
nu0     = 6
tau20   = 20/3
iV0     = diag(1/diag(V0))
n0sig202= n0*sig20/2
c((n0*sig20/2)/(n0/2-1),((n0*sig20/2)/(n0/2-1))^2/(n0/2-2),(nu0*tau20/2)/(nu0/2-1),((nu0*tau20/2)/(nu0/2-1))^2/(nu0/2-2))

# MCMC setup
set.seed(122545)
v  = 0.1
M0 = 100000
M  = 1000
L  = 	50
niter = M0+M*L
theta = ptr[1:3]
tau2  = ptr[4]
sig2  = ptr[5]
x     = xtr
xs    = NULL
ps    = NULL
n0n2  = (n0+n)/2
runtime = system.time({
for (iter in 1:niter){	
  print(iter)
  x     = draw.x(y,x,x0,theta,sig,tau,v)
  sig2  = 1/rgamma(1,n0n2,n0sig202+sum((y-x^2/20)^2)/2)
  sig   = sqrt(sig2)
  W     = cbind(c(x0,x[1:(n-1)]),0:(n-1))
  Z     = t(apply(W,1,g))
  par   = fixedpar(x,Z,theta0,iV0,nu0,tau20)
  theta = par[1:3]
  tau2  = par[4]
  tau   = sqrt(tau2)
  if (iter>M0){
    xs    = rbind(xs,x)
    ps    = rbind(ps,c(par,sig2))
  }
}
})
index = seq(1,niter-M0,by=L) 
length(index)

mx = apply(xs[index,],2,median)
lx = apply(xs[index,],2,quant025)
ux = apply(xs[index,],2,quant975)

pdf(file="nonlinear-x-longer.pdf",width=10,height=7)
par(mfrow=c(1,1))
plot(xtr,type="l",ylim=range(lx,ux,xtr),xlab="time",ylab="")
lines(lx,col=2)
lines(ux,col=2)
lines(mx,col=4)
lines(xtr,col=1)
legend(0,40,legend=c("TRUE","POSTERIOR MEAN","90%CI"),col=c(1,4,2),lty=rep(1,3),lwd=rep(1,3),bty="n")
dev.off()

names = c("alpha","beta","gamma","tau2","sig2")
pdf(file="nonlinear-par-longer.pdf",width=15,height=10)
par(mfrow=c(3,5))
for (i in 1:5){
  plot(index,ps[index,i],type="l",xlab="iteration",ylab="",main=names[i])
  abline(h=ptr[i],col=2,lwd=4)
}
for (i in 1:5)
  acf(ps[index,i],main="")
for (i in 1:5){
  hist(ps[index,i],xlab="",ylab="",main="")
  abline(v=ptr[i],col=2,lwd=4)
}
dev.off()

acfs=NULL
for (t in 1:n)
  acfs = rbind(acfs,acf(xs[index,t],plot=FALSE)$acf)
pdf(file="nonlinear-acf-longer.pdf",width=10,height=7)
par(mfrow=c(1,1))
ts.plot(t(acfs),xlab="lag",ylab="ACF",col=gray(0.8))
dev.off()