############################################################################################
# 
# AR(1) PLUS NOISE MODEL
#
############################################################################################
#
# 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@ChicagoGSB.edu
# URL: http://faculty.chicagobooth.edu/hedibert.lopes/research/
#
############################################################################################
rm(list=ls())

PL = function(y,tau2s,sig2s,xs){
  n      = length(y)
  N      = length(xs)
  draws  = array(0,c(n,N,3))
  s      = matrix(0,N,4)
  s[,1]  = n0
  s[,2]  = n0*sig20
  s[,3]  = nu0
  s[,4]  = nu0*tau20 
  for (t in 1:n){
   	 print(t)
     # Resampling
     stdevs = sqrt(sig2s+tau2s)
     weight = dnorm(y[t],xs,stdevs)
  	   k      = sample(1:N,size=N,replace=T,prob=weight)
  	   tau2s  = tau2s[k]
  	   sig2s  = sig2s[k]
  	   xs1    = xs[k]
  	   so     = s[k,]
    	 # Propagating
  	   vars   = 1/(1/sig2s+1/tau2s)
  	   sds    = sqrt(vars)
  	   means  = vars*(y[t]/sig2s + xs1/tau2s)
  	   xs     = rnorm(N,means,sds)
  	  	 # Updating sufficient stats
  	   s[,1] = so[,1] + 1
  	   s[,2] = so[,2] + (y[t]-xs)^2
  	   s[,3] = so[,3] + 1
     s[,4] = so[,4] + (xs-xs1)^2
     # Sampling parameters
     sig2s  = 1/rgamma(N,s[,1]/2,s[,2]/2)
     tau2s  = 1/rgamma(N,s[,3]/2,s[,4]/2)
     # Storing quantiles
     draws[t,,] = cbind(tau2s,sig2s,xs)
  }
  return(draws)
}


data = read.table("SA-marketindex.txt",header=TRUE)
n    = nrow(data)
y    = data[,2]
ind  = seq(1,n,length=5)
dat  = data[ind,1]

pdf(file="dlm.pdf",width=20,height=10)

#pdf(file="data.pdf",width=10,height=8)
par(mfrow=c(1,1))
plot(y,ylab="",main="Prices",axes=FALSE,xlab="Days",type="l")
axis(2);axis(1,at=ind,lab=dat);box()
#dev.off()

# Prior hyperparameters
m0    = 50
C0    = 100
n0    = 4.1
sig20 = 1
nu0   = 4.1
tau20 = 0.5

# Particle learning
set.seed(246521)
N      = 10000
xs     = rnorm(N,m0,sqrt(C0))
tau2s  = 1/rgamma(N,nu0/2,nu0*tau20/2)
sig2s  = 1/rgamma(N,n0/2,nu0*sig20/2)
pl     = PL(y,tau2s,sig2s,xs)

qpl = array(0,c(n,3,3))
for (i in 1:3)
  qpl[,i,] = t(apply(pl[,,i],1,quantile,c(0.05,0.5,0.95)))

#pdf(file="dlm-graph1.pdf",width=10,height=8)
par(mfrow=c(1,1))
ind1 = 1:n
plot(qpl[,1,1],main="",ylim=range(qpl[ind1,1:2,]),lwd=2,axes=FALSE,xlab="Days",type="l",ylab="")
axis(2);axis(1,at=ind,lab=dat);box()
lines(qpl[,1,2],lwd=2);lines(qpl[,1,3],lwd=2)
for (i in 1:3) lines(qpl[,2,i],lwd=2,col=2)
legend(n/2,5,legend=c(expression(sigma^2),expression(tau^2)),bty="n",col=1:2,lwd=c(2,2),cex=2)
#dev.off()

ind2 = trunc(seq(1,n,length=16))
ind3 = 30:n
ind4 = sample(1:N,size=N/10,replace=FALSE,prob=rep(1/N,N))
Lx = min(pl[ind3,,1])
Ux = max(pl[ind3,,1])
Ly = min(pl[ind3,,2])
Uy = max(pl[ind3,,2])
pdf(file="dlm-graph2.pdf",width=12,height=12)
par(mfrow=c(4,4))
for (j in 1:16){
  plot(pl[ind2[j],ind4,1:2],xlim=c(Lx,Ux),ylim=c(Ly,Uy),xlab=expression(sigma^2),
       ylab=expression(tau^2),main="",pch=".")
  title(data[ind2[j],1])
}
dev.off()

#pdf(file="dlm-graph3.pdf",width=20,height=10)
par(mfrow=c(1,1))
plot(y,main="",lwd=2,axes=FALSE,xlab="Days",pch=16,ylab="",cex=0.5)
axis(2);axis(1,at=ind,lab=dat);box()
for (i in 1:n)
  segments(i,qpl[i,3,1],i,qpl[i,3,3],col=2) 
points(y,pch=16,cex=0.5)
#dev.off()

dev.off()