###################################################################
#
#
#  two factor Stochastic Volatility model
#
#
###################################################################
#
# Author : Hao Liya
#
###################################################################
## simulation
set.seed(1243)
gamma     = -1.5
G         = matrix(c(0.98,0,0,0.6),2,2)
W1         = 0.2^2
W2         = 0.8^2
lambda0   = c(0,0)
n         = 1000
lambda    = matrix(0,n,2)
slam      = rep(0,n)
error1     = rnorm(n,0.0,sqrt(W1))
error2     = rnorm(n,0.0,sqrt(W2))
error      = array(c(error1,error2),c(n,2))
lambda[1,] = G%*%lambda0+error[1,]
slam[1]    = lambda[1,1]+lambda[1,2]
for (t in 2:n){
  lambda[t,] = G%*%lambda[t-1,]+error[t,]
  slam[t]=lambda[t,1]+lambda[t,2]
}
h = exp((slam+gamma)/2)
y = rnorm(n,0,h)
plot(y,type="l")

pars0      = c(gamma,G[1,1],G[2,2],sqrt(W1),sqrt(W2))


## functions
ffbsu = function(y,alpha,V,G,gamma,W,a1,R1){
  nd= 2
  n = length(y)
  F = matrix(c(1,1),2,1)
  a = matrix(0,n,2)
  R = matrix(0,2,2)
  m = matrix(0,n,2)
  C = array(0,c(n,2,2))
  B = array(0,c(n-1,2,2))
  H = array(0,c(n-1,2,2))
  pH= array(0,c(n-1,2,2))
  # time t=1
  a[1,]= a1
  R    = R1
  f    = alpha[1]+t(F)%*%a[1,]+gamma
  Q    = t(F)%*%R%*%F+V[1]
  A    = R%*%F/Q[1,1]
  m[1,] = a[1,]+A%*%(y[1]-f)
  C[1,,] = R-A%*%t(A)*Q[1,1]
  # forward filtering
  for (t in 2:n){
    a[t,]= G%*%m[t-1,]
    R    = G%*%C[t-1,,]%*%t(G) + W
    f    = gamma+alpha[t]+t(F)%*%a[t,]
    Q    = t(F)%*%R%*%F+V[t]
    A    = R%*%F/Q[1,1]
    m[t,] = a[t,]+A%*%(y[t]-f)
    C[t,,] = R-A%*%t(A)*Q[1,1]
    B[t-1,,] = C[t-1,,]%*%t(G)%*%solve(R)
    H[t-1,,] = C[t-1,,]-B[t-1,,]%*%G%*%C[t-1,,]

  }
  # backward sampling
  theta = matrix(0,nd,n)
  eta   = matrix(0,nd,n)
  eta[1,]=alpha
  eta[2,]=V

  theta[,n] = rmvnorm(1,m[n,],C[n,,])
  for (t in (n-1):1)
    theta[,t] = rmvnorm(1,m[t,]+B[t,,]%*%(theta[,t+1]-a[t+1,]),H[t,,])
  return(theta,eta)
}

#-------------------------------------------------------
#-------------------------------------------------------
fmu  = function(y,slam,peta,m,v){
   a = peta[1,]
   b = peta[2,]
   v1= 1/(1/v+sum(1/b))
   m1= v1*sum((log(y^2)-slam-a)/b)+m/(1+v*sum(1/b))
   gamma=rnorm(1,m1,sqrt(v1))
   return(gamma,m1,v1)
}

fphi = function(y,plambda,pW1,pW2){
   n  =length(y)
   m1 = sum(plambda[1,2:n]*plambda[1,1:(n-1)])/sum(plambda[1,1:(n-1)]^2)
   m2 = sum(plambda[2,2:n]*plambda[2,1:(n-1)])/sum(plambda[2,1:(n-1)]^2)
   v1 = pW1/sum(plambda[1,1:(n-1)]^2)
   v2 = pW2/sum(plambda[2,1:(n-1)]^2)
   phi2 = rnorm(1,m2,sqrt(v2))
   repeat{
   phi1 = rnorm(1,m1,sqrt(v1))
   if(phi1>phi2)
   break
   }
   return(phi1,phi2)
}

fsig=function(alpha1,alpha2,beta1,beta2,pG,plambda,y){
   n = length(y)
   palpha1 = (n+2*alpha1)/2
   palpha2 = (n+2*alpha2)/2
   pbeta1  = 2*beta1+plambda[1,1]^2*(1-pG[1,1]^2)+sum((plambda[1,2:n]-pG[1,1]*plambda[1,1:(n-1)])^2)
   pbeta1  = pbeta1/2
   pbeta2  = 2*beta2+plambda[2,1]^2*(1-pG[2,2]^2)+sum((plambda[2,2:n]-pG[2,2]*plambda[2,1:(n-1)])^2)
   pbeta2  = pbeta2/2
   psig1   = rinvgamma(1, palpha1, pbeta1)
   psig2   = rinvgamma(1, palpha2, pbeta2)
   return(psig1,psig2,palpha1,palpha2,pbeta1,pbeta2)
}

#----------------------------------------------------------------------------
# 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)))
}

#------------------------------------------------
# Quantile statistics
#------------------------------------------------
quant05 = function(x){quantile(x,0.05)}
quant95 = function(x){quantile(x,0.95)}

#------------------------------------------------
# Sampling the log-volatilities in a
# standard univariate stochastic volatility model
#------------------------------------------------
svu = function(y,pslam,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-pslam-gamma,ncind,mu,sig,q)
  ffbsu(y,mu[z],sig2[z],G,gamma,W,a1,R1)
}


##############################################################################
#  STOCHASTIC VOLATILITY + FORWARD FILTERING BACKWARD SAMPLING               #
##############################################################################

# Prior hyperparameters
m         = 0
v         = 1000
alpha1    = 1.5
alpha2    = 1
beta1     = 0.055
beta2     = 0.5
a1        = c(0,0)
R1        = diag(1000,2)

# Initial values

pgamma     = 0
pG         = matrix(c(0.95,0,0,0.5),2,2)
pW1        =0.1^2
pW2        =0.9^2
pW         =diag(c(pW1,pW2))
perror1    = rnorm(n,0.0,sqrt(pW1))
perror2    = rnorm(n,0.0,sqrt(pW2))
perror     = array(c(perror1,perror2),c(n,2))
plambda    = matrix(0,2,n)
pslam      =rep(0,n)
plambda[,1] = pG%*%lambda0+perror[1,]
pslam[1]    = plambda[1,1]+plambda[2,1]
for (t in 2:n){
  plambda[,t] = pG%*%plambda[,t-1]+perror[t,]
  pslam[t]=plambda[1,t]+plambda[2,t]
}
lambdas   = plambda
pars      = c(pgamma,pG[1,1],pG[2,2],sqrt(pW1),sqrt(pW2))
M0        = 1000
M         = 1000

date()
for (iter in 1:(M0+M)){
  pl       = svu(y,pslam,pgamma,pG,pW,a1,R1)
  plambda  = pl$theta
  peta     = pl$eta
  F        = c(1,1)
  pslam    = F%*%plambda
  param1   = fmu(y,pslam,peta,m,v)
  pgamma   = param1$gamma
  m        = param1$m1
  v        = param1$v1
  param2   = fphi(y,plambda,pW1,pW2)
  pG       = diag(c(param2$phi1,param2$phi2))
  param3   = fsig(alpha1,alpha2,beta1,beta2,pG,plambda,y)
  pW1      = param3$psig1
  pW2      = param3$psig2
  pW       = diag(c(pW1,pW2))
  alpha1   = param3$palpha1
  alpha2   = param3$palpha2
  beta1    = param3$pbeta1
  beta2    = param3$pbeta2
  pars    = rbind(pars,c(pgamma,pG[1,1],pG[2,2],sqrt(pW1),sqrt(pW2)))
  lambdas = rbind(lambdas,plambda)
}
date()

## description
odd   = rep(0,(M0+M+1))
even  = rep(0,(M0+M+1))
for(t in 1:(M0+M+1)){
odd[t] = 2*t-1
even[t]= 2*t
}
lambdas1 = lambdas[odd,]
lambdas2 = lambdas[even,]

mlambdas1  = apply(lambdas1[(M0+1):(M0+M),],2,mean)
mlambdas2  = apply(lambdas2[(M0+1):(M0+M),],2,mean)
par(mfrow=c(3,1))
ts.plot(lambda[,1])
lines(mlambdas1+1.5,col=2)
ts.plot(lambda[,2])
lines(mlambdas2,col=2)
ts.plot(lambda[,1]+lambda[,2])
lines(mlambdas1+mlambdas2+1.5,col=2)

x11()
h105 = apply(lambdas1[(M0+1):(M0+M),],2,quant05)
h205 = apply(lambdas2[(M0+1):(M0+M),],2,quant05)
h195 = apply(lambdas1[(M0+1):(M0+M),],2,quant95)
h295 = apply(lambdas2[(M0+1):(M0+M),],2,quant95)
par(mfrow=c(3,1))
plot(lambda[,1],type="l",ylim=c(min(h105),max(h195)+1.5),xlab="time",ylab="h1")
lines(lambda[,1],col=1,lwd=1)
lines(mlambdas1+1.5,col=2,lwd=2)
lines(h105+1.5,col=3,lwd=2)
lines(h195+1.5,col=3,lwd=2)
plot(lambda[,2],type="l",ylim=c(min(h205),max(h295)),xlab="time",ylab="h2")
lines(lambda[,2],col=1,lwd=1)
lines(mlambdas2,col=2,lwd=2)
lines(h205,col=3,lwd=2)
lines(h295,col=3,lwd=2)
plot(lambda[,1]+lambda[,2],type="l",ylim=c(min(h105+h205),max(h195+h295)+1.5),xlab="time",ylab="h1+h2")
lines(lambda[,1]+lambda[,2],col=1,lwd=1)
lines(mlambdas1+mlambdas2+1.5,col=2,lwd=2)
lines(h105+h205+1.5,col=3,lwd=2)
lines(h195+h295+1.5,col=3,lwd=2)



x11()
names = c("mu","phi1","phi2","sigma1","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=pars0[i],col=2)
  hist(pars[(M0+1):(M0+M),i],xlab="",ylab="",main="")
  abline(v=pars0[i],col=2)
}
par(mfrow=c(3,2))
for (i in 4:5){
  plot(pars[(M0+1):(M0+M),i],type="l",xlab="iteration",ylab="",main=names[i])
  abline(h=pars0[i],col=2)
  hist(pars[(M0+1):(M0+M),i],xlab="",ylab="",main="")
  abline(v=pars0[i],col=2)
}


x11()
hs  = exp(lambdas1[(M0+1):(M0+M),]+lambdas2[(M0+1):(M0+M),]-1.5)
mh  = apply(hs,2,mean)
h05 = apply(hs,2,quant05)
h95 = apply(hs,2,quant95)
h   = exp(lambda[,1]+lambda[,2]-1.5)
par(mfrow=c(1,1))
plot(3+h,type="l",ylim=c(0.0,10),xlab="time",ylab="")
lines(3+h,col=1,lwd=1)
lines(3+mh,col=2,lwd=2)

legend(770,10,legend=c("h(t)","E(h(t)|data)","data"),col=c(1,2,3),lty=rep(1,3),lwd=rep(2,3),bty="n",cex=1)
lines(abs(y)/max(abs(y))*4,type="h",col=3)