#-------------------------------------------------------
# Univariate FFBS: 
# y(t)     ~ N(alpha(t)+F(t)*theta(t);V(t))        
# theta(t) ~ N(gamma+G*theta(t-1);W)               
#-------------------------------------------------------
ffbsu = function(y,F,alpha,V,G,gamma,W,a1,R1,nd=1){
  n = length(y)
  if (length(F)==1){F = rep(F,n)}
  if (length(alpha)==1){alpha = rep(alpha,n)}
  if (length(V)==1){V = rep(V,n)}
  a = rep(0,n)
  R = rep(0,n)
  m = rep(0,n)
  C = rep(0,n)
  B = rep(0,n-1)
  H = rep(0,n-1)
  # time t=1
  a[1] = a1
  R[1] = R1
  f    = alpha[1]+F[1]*a[1]
  Q    = R[1]*F[1]**2+V[1]
  A    = R[1]*F[1]/Q
  m[1] = a[1]+A*(y[1]-f)
  C[1] = R[1]-Q*A**2
  # forward filtering
  for (t in 2:n){
    a[t] = gamma + G*m[t-1]
    R[t] = C[t-1]*G**2 + W
    f    = alpha[t]+F[t]*a[t]
    Q    = R[t]*F[t]**2+V[t]
    A    = R[t]*F[t]/Q
    m[t] = a[t]+A*(y[t]-f)
    C[t] = R[t]-Q*A**2
    B[t-1] = C[t-1]*G/R[t]
    H[t-1] = sqrt(C[t-1]-R[t]*B[t-1]**2)
  }
  # backward sampling
  theta = matrix(0,nd,n)
  theta[,n] = rnorm(nd,m[n],sqrt(C[n]))
  for (t in (n-1):1)
    theta[,t] = rnorm(nd,m[t]+B[t]*(theta[,t+1]-a[t+1]),H[t])
  if (nd==1){
    theta[1,]
  }
  else{
    theta
  }
}

#-----------------------------------
# y|X,beta,sig2 ~ N(X*beta,sig2*I_n)
# beta|sig2     ~ N(b,sig2*inv(A))
#      sig2     ~ IG(v/2,v*lam/2)
#-----------------------------------
fixedpar = function(y,X,A,b,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(2))
  return(c(mean,sig2))
}

#---------------------------------------------
# Sample Z from 1,2,...,k, with 
# P(Z=i) proportional to q[i]N(mu[i],sig2[i])
#---------------------------------------------
ncind = function(y,mu,sig,q){
  w = dnorm(y,mu,sig)*q
  return(sample(1:length(q),size=1,prob=w/sum(w)))
}

#------------------------------------------------
# Sampling the log-volatilities in a 
# standard univariate stochastic volatility model
#------------------------------------------------
svu = function(y,lambda,gamma,G,W,a1,R1){
  mu   = c(-11.40039,-5.24321,-9.83726,1.50746,-0.65098,0.52478,-2.35859)
  sig2 = c(5.795960,2.613690,5.179500,0.167350,0.640090,0.340230,1.262610)
  q    = c(0.007300,0.105560,0.000020,0.043950,0.340010,0.245660,0.257500)
  y    = log(y^2)
  sig  = sqrt(sig2)
  z    = sapply(y-lambda,ncind,mu,sig,q)
  ffbsu(y,1.0,mu[z],sig2[z],G,gamma,W,a1,R1)
}


######################################
#  FFBS for the SV-AR(1) model
######################################
set.seed(1243)
gamma     = -0.00645
G         = 0.99
W         = 0.15^2
lambda0   = 0.0
n         = 1000
lambda    = rep(0,n)
error     = rnorm(n,0.0,sqrt(W))
lambda[1] = gamma+G*lambda0+error[1]
for (t in 2:n)
  lambda[t] = gamma+G*lambda[t-1]+error[t]
h = exp(lambda/2)
y = rnorm(n,0,h)
plot(y,type="l")

# Prior hyperparameters
b         = c(0,0)
A         = diag(0.00000001,2)
v         = 0.00000002
lam       = 1.00000000
a1        = 0
R1        = 1000

# Initial values
lambdas   = lambda
pars      = c(gamma,G,sqrt(W))
M0        = 1000
M         = 1000

system.time({
for (iter in 1:(M0+M)){
  lambda  = svu(y,lambda,gamma,G,W,a1,R1)
  X       = cbind(1,lambda[2:n])
  param   = fixedpar(lambda[1:(n-1)],X,A,b,v,lam)
  gamma   = param[1]
  G       = param[2]
  W       = param[3]
  pars    = rbind(pars,c(param[1:2],sqrt(param[3])))
  lambdas = rbind(lambdas,lambda)
}
})

hs  = exp(lambdas[(M0+1):(M0+M),])
mh  = apply(hs,2,mean)
h05 = apply(hs,2,quantile,0.05)
h95 = apply(hs,2,quantile,0.95)

names = c("mu","phi","sigma2")
par(mfrow=c(3,2))
for (i in 1:3){
  plot(pars[(M0+1):(M0+M),i],type="l",xlab="iteration",ylab="",main=names[i])
  abline(h=pars[1,i],col=2)
  hist(pars[(M0+1):(M0+M),i],xlab="",ylab="",main="")
  abline(v=pars[1,i],col=2)
}

par(mfrow=c(1,1))
plot(1+h,type="l",ylim=c(0.0,max(1+h95m)),xlab="time",ylab="")
lines(1+h,col=1,lwd=1)
lines(1+mh,col=2,lwd=1)
lines(1+h05,col=3,lwd=1)
lines(1+h95,col=3,lwd=1)
legend(400,4,legend=c("h(t)","E(h(t)|data)","90%CI"),col=c(1,2,3),lty=rep(1,3),lwd=rep(1,3),bty="n")
lines(abs(y)/max(abs(y)),type="h")

install.package("stochvol")
library("stochvol")

fit = svsample(y)

ind = (M0+1):(M0+M)
par(mfrow=c(2,2))
plot(density(pars[ind,1]),xlab="",main=expression(beta[0]))
lines(density(fit$para[,1]*(1-fit$para[,2])),col=2)
legend("topleft",legend=c("Naive coding","Optimized coding"),col=1:2,lty=1)
plot(density(pars[ind,2]),xlab="",main=expression(beta[1]))
lines(density(fit$para[,2]),col=2)
plot(density(pars[ind,3]),xlab="",main=expression(tau^2))
lines(density(fit$para[,3]),col=2)
ts.plot(mh,ylim=c(0,3))
lines(apply(exp(fit$latent),2,median),col=2)