#################################################################################################
#
# FIRST ORDER DLM - PARTICLE LEARNING
#
#  y(t) = x(t)   + e(t)   e(t) ~ N(0,sig2)
#  x(t) = x(t-1) + w(t)   w(t) ~ N(0,tau2)
#  
#  x(0) ~ N(m0,V0) and sig2 ~ IG(a0,b0)
#
#################################################################################################
#
# 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())
q025       = function(x){quantile(x,0.025)}
q975       = function(x){quantile(x,0.975)}
dinvgamma  = function(x,a,b){dgamma(1/x,a,b)*(1/x^2)}
drawnorm   = function(x){rmvnorm(1,x[1:2],matrix(x[3:6],2,2))}
likelihood = function(x,mu,sig){dnorm(x,mu,sig)}

ffbslike = function(y,sig2,tau2,d0,V0){
  n=length(y);like=0.0
  for (t in 1:n){
    if (t==1){a=d0;R=V0+tau2}else{a=m;R=C+tau2}
    f=a;Q=R+sig2;A=R/Q;m=a+A*(y[t]-f);C=R-Q*A**2
    like=like+dnorm(y[t],f,sqrt(Q),log=T)
  }
  return(like)
}

PL = function(y,tau2,a0,b0,m0,C0,M){
  n     = length(y)
  x     = rnorm(M,m0,sqrt(C0))
  sig2  = 1/rgamma(M,a0,b0)
  S     = cbind(rep(a0,M),rep(b0,M))
  sig2s = NULL
  for (t in 0:(n-1)){
  	 A     = tau2/(tau2+sig2)
    w     = dnorm(y[t+1],x,sqrt(sig2+tau2))
    k     = sample(1:M,size=M,prob=w,replace=T)
    var   = A[k]*sig2[k]
    mean  = A[k]*y[t+1]+(1-A[k])*x[k]
    x     = rnorm(M,mean,sqrt(var))
    S[,1] = S[k,1] + 1/2
    S[,2] = S[k,2] + (y[t+1]-x)^2/2
    sig2  = 1/rgamma(M,S[,1],S[,2])
    sig2s = rbind(sig2s,sig2)
  }
  return(sig2s)
}

########################################################################################################
pdf(file="pl-tau2-n1000.pdf",width=11,height=3.5)
par(mfrow=c(1,3))
n     = 1000
sig2  = 0.1
tau2s = c(0.05,0.1,0.2)
sig   = sqrt(sig2)
for (tau2 in tau2s){
  # Data simulation
  # ---------------
  set.seed(12345)
  tau   = sqrt(tau2)
  x     = rep(0,n+1)
  y     = rep(0,n+1)
  x[1]  = 25
  y[1]  = rnorm(1,x[1],sig)
  for (t in 2:(n+1)){
    x[t] = rnorm(1,x[t-1],tau)
    y[t] = rnorm(1,x[t],sig)
  }
  x0 = x[1]
  x  = x[2:(n+1)]
  y  = y[2:(n+1)]

  # True values
  x0true   = x0
  xtrue    = x
  sig2true = sig2
  tau2true = tau2

  # Prior setup
  # -----------
  a0 = 5
  b0 = (a0-1)*sig2
  m0 = x0true
  C0 = 100

  # True sequential posterior for sig2
  # ----------------------------------
  L      = 0.0001
  U      = 5*sig2true
  N      = 300
  sig2s  = seq(L,U,length=N)
  hsig2  = sig2s[2]-sig2s[1]
  ns     = seq(50,n,by=50)
  qs     = NULL
  for (i in ns){
    logpost = rep(0,N)
    for (j in 1:N)
      logpost[j] = ffbslike(y[1:i],sig2s[j],tau2true,m0,C0)+log(dinvgamma(sig2s[j],a0,b0))
    postsig2 = exp(logpost-max(logpost))
    postsig2 = postsig2/sum(postsig2)
    postsig2 = postsig2/hsig2
    q        = cumsum(postsig2*hsig2)
    quant    = c(max(sig2s[q<0.025]),max(sig2s[q<0.5]),max(sig2s[q<0.975]))
    qs       = rbind(qs,quant)
  }

  # Particle Learning (PL)
  # ----------------------
  set.seed(54321)
  M   = 2000
  pl  = PL(y,tau2,a0,b0,m0,C0,M)
  mpl = apply(pl,1,median)
  lpl = apply(pl,1,q025)
  upl = apply(pl,1,q975)

  L = min(lpl[ns],qs)
  U = max(lpl[ns],qs)
  plot(ns,mpl[ns],type="l",xlab="time",ylab="",lwd=1,ylim=c(L,U),col=2,main="")
  title(paste("tau/sig=",round(sqrt(tau2/sig2),2),sep=""))
  lines(ns,lpl[ns],col=2,lwd=1,lty=1)
  lines(ns,upl[ns],col=2,lwd=1,lty=1)
  lines(ns,qs[,1],col=4)
  lines(ns,qs[,2],col=4)
  lines(ns,qs[,3],col=4)
  abline(h=sig2true,lwd=1,col=3)
}
dev.off()