#install.packages("stochvol")
#library("stochvol")

mu    = -10
phi   = 0.99
sigma = 0.2
n     = 1000

sim   = svsim(n,mu=mu,phi=phi,sigma=sigma)

pdf(file="sv-mcmc-smc.pdf",width=12,height=9)
par(mfrow=c(1,1))
ts.plot(sim$y)
lines(sim$vol,col=2)
lines(-sim$vol,col=2)

run   = svsample(sim$y)
quant = t(apply(exp(run$latent/2),2,quantile,c(0.05,0.5,0.95)))

par(mfrow=c(1,1))
ts.plot(quant,ylim=range(quant,stdev))
lines(sim$vol,col=2)

par(mfrow=c(2,3))
ts.plot(run$para[,1],ylab="",xlab="iterations",main=expression(mu))
abline(h=mu,col=2)
ts.plot(run$para[,2],ylab="",xlab="iterations",main=expression(phi))
abline(h=phi,col=2)
ts.plot(run$para[,3],ylab="",xlab="iterations",main=expression(sigma))
abline(h=sigma,col=2)
acf(run$para[,1],main="")
acf(run$para[,2],main="")
acf(run$para[,3],main="")

ts = seq(20,n,by=20)
nt = length(ts)
draws = array(0,c(nt,10000,4))
for (i in 1:nt){
  run   = svsample(sim$y[1:ts[i]])
  draws[i,,1:3] = run$para
  draws[i,,4] = exp(run$latent[,ts[i]]/2)
}

qs = array(0,c(nt,3,4))
for (i in 1:4)
  qs[,,i] = t(apply(draws[,,i],1,quantile,c(0.05,0.5,0.95)))

names = c("mu","phi","sigma")
par(mfrow=c(1,3))
for (i in 1:3){
  plot(ts,qs[,2,i],ylim=range(qs[,,i]),xlab="Time",
       type="l",ylab="",main=names[i])
  lines(ts,qs[,1,i],lty=2)
  lines(ts,qs[,3,i],lty=2)
  points(ts[nt],quantile(run$para[,i],0.05),col=2,pch=16)
  points(ts[nt],quantile(run$para[,i],0.5),col=2,pch=16)
  points(ts[nt],quantile(run$para[,i],0.95),col=2,pch=16)
}
legend("topright",legend=c("On-line","Off-line"),col=1:2,lty=1)

par(mfrow=c(1,1))
plot(ts,qs[,2,4],ylim=range(qs[,,4]),type="l",xlab="Time",
     ylab="",main="stdev(t)")
lines(ts,qs[,1,4],lty=2)
lines(ts,qs[,3,4],lty=2)
lines(ts,quant[ts,1],lty=2,col=2)
lines(ts,quant[ts,2],col=2)
lines(ts,quant[ts,3],lty=2,col=2)
legend("topright",legend=c("On-line","Off-line"),col=1:2,lty=1)

time = c(0.417,0.518,0.556,0.638,0.716,0.732,0.871,0.922,0.984,1.052,
1.149,1.193,1.3,1.687,1.439,1.53,1.593,1.672,1.741,1.791,
1.847,1.927,2.028,2.144,2.166,2.476,2.337,2.377,2.496,2.549,
2.589,2.679,2.748,2.981,3.155,2.957,3.062,3.263,3.638,3.219,
3.272,3.398,4.157,4.184,3.684,3.843,3.718,3.77,3.893,4.035)

plot(ts,time,xlab="Number of observations",ylab="Elapsed time")
title(paste("total time =",sum(time)," seconds"))

# Sequential learning
M = 10000
hss = matrix(0,M,n)
hs = rnorm(M,mu,sigma)
for (t in 1:n){
  hs1 = rnorm(M,mu+phi*(hs-mu),sigma)
  w = dnorm(sim$y[t],0,exp(hs1/2))
  k = sample(1:M,size=M,replace=TRUE,prob=w)
  hs = hs1[k]
  hss[,t] = hs
}

qmcmc = t(apply(exp(hss/2),2,quantile,c(0.05,0.5,0.95)))

ts.plot(qmcmc)
points(ts,qs[,1,4],col=2,pch=16,cex=0.5)
points(ts,qs[,2,4],col=2,pch=16,cex=0.5)
points(ts,qs[,3,4],col=2,pch=16,cex=0.5)
legend("topright",legend=c("On-line (SMC)","Online (brute force MCMC)"),col=1:2,lty=1)

ts.plot(qmcmc)
lines(quant[,1],col=2)
lines(quant[,2],col=2)
lines(quant[,3],col=2)
legend("topright",legend=c("On-line (SMC)","Off-line (MCMC)"),col=1:2,lty=1)


dev.off()