############################################################################################
# 
# FIRST ORDER DYNAMIC LINEAR 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())
DLM.SIM = function(n,tau,sig,x0){
  w     = rnorm(n,0,tau)
  v     = rnorm(n,0,sig)
  x     = rep(0,n)
  y     = rep(0,n)
  x[1]  = alpha+beta*x0 + w[1]
  y[1]  = x[1] + v[1]
  for (t in 2:n){
    x[t] = alpha+beta*x[t-1] + w[t]
    y[t] = x[t] + v[t]
  }
  return(list(x=x,y=y))
}

DLM = function(y,alpha,beta,sig2,tau2,m0,C0){
  n = length(y)
  m = m0
  C = C0
  l = 0.0
  mm = rep(0,n)
  CC = rep(0,n)
  for (t in 1:n){
    a = alpha+beta*m
    R = beta^2*C+tau2
    Q = R+sig2
    l = l+dnorm(y[t],a,sqrt(Q),log=TRUE)
    A = R/Q
    m = (1-A)*a+A*y[t]
    C = A*sig2
    mm[t] = m
    CC[t] = C
 }
 return(list(m=mm,C=CC,l=l))
} 

SIS = function(y,alpha,beta,sig2,tau2,m0,C0,M){
  n  = length(y)
  x  = rnorm(M,m0,sqrt(C0))
  qs = matrix(0,n,3)
  neff = rep(0,n)
  mx = rep(0,n)
  w  = rep(1/M,M)
  for (t in 1:n){
  	  x      = rnorm(M,alpha+beta*x,tau)
    w1      = w/sum(w)
  	  w      = w*dnorm(y[t],x,sig)
    k       = sample(1:M,size=M,prob=w,replace=T)
    x1      = x[k]
    qs[t,]  = quantile(x1,c(0.025,0.5,0.975)) 
    neff[t] = 1/sum(w1^2)
    mx[t]   = sum(w1*x)
  }
  return(list(qs=qs,neff=neff,mx=mx))
}

SIS1 = function(y,alpha,beta,sig2,tau2,m0,C0,M,perc){
  n  = length(y)
  x  = rnorm(M,m0,sqrt(C0))
  qs = matrix(0,n,3)
  neff = rep(0,n)
  mx = rep(0,n)
  w  = rep(1/M,M)
  M1 = perc*M
  for (t in 1:n){
  	  x      = rnorm(M,alpha+beta*x,tau)
  	  w      = w*dnorm(y[t],x,sig)
    k       = sample(1:M,size=M,prob=w,replace=T)
    x1      = x[k]
    qs[t,]  = quantile(x1,c(0.025,0.5,0.975)) 
    w1      = w/sum(w)
    neff[t] = 1/sum(w1^2)
    mx[t] = sum(x*w1)
    if (neff[t]<M1){
      x = x1
      w = rep(1/M,M)
      mx[t] = mean(x)
    }
  }
  return(list(qs=qs,neff=neff,mx=mx))
}

BF = function(y,alpha,beta,sig2,tau2,m0,C0,M){
  n  = length(y)
  x  = rnorm(M,m0,sqrt(C0))
  qs = matrix(0,n,3)
  neff = rep(0,n)
  mx   = rep(0,n)
  for (t in 1:n){
  	  x      = rnorm(M,alpha+beta*x,tau)
  	  w      = dnorm(y[t],x,sig)
    k      = sample(1:M,size=M,prob=w,replace=T)
    x      = x[k]
    w1      = w/sum(w)
    mx[t]   = sum(w1*x)
    neff[t] = 1/sum(w1^2)
    qs[t,]  = quantile(x,c(0.025,0.5,0.975)) 
  }
  return(list(qs=qs,neff=neff,mx=mx))
}

# Simulating 1st order DLM
# ------------------------
M     = 1000
n     = 50
sig2s = c(0.25,0.5,1.0)
alpha = 0.0
beta  = 1.0
tau2  = 0.2
x0    = 0.0
m0    = 0.0
C0    = 10

set.seed(12345)
y = matrix(0,n,3)
x = matrix(0,n,3)
qdlm = array(0,c(3,n,3))
pdf(file="rmae1.pdf",width=15,height=10)
par(mfrow=c(2,3))
for (i in 1:3){
  sig2  = sig2s[i]
  sig   = sqrt(sig2)
  tau   = sqrt(tau2)
  sim   = DLM.SIM(n,tau,sig,x0)
  y[,i] = sim$y
  x[,i] = sim$x
  plot(y[,i],type="b",main=paste("sig2=",sig2,sep=""),xlab="time",pch=16)
  lines(x[,i],col=2,pch=16,type="b",lwd=2)
}
for (i in 1:3){
  dlm  = DLM(y[,i],alpha,beta,sig2,tau2,m0,C0)
  qdlm[i,,] = cbind(dlm$m+qnorm(0.025)*sqrt(dlm$C),dlm$m,dlm$m+qnorm(0.975)*sqrt(dlm$C))
  ts.plot(qdlm[i,,])
  points(x[,i],col=2,pch=16,lwd=2)
}
dev.off()


pdf(file="rmae2.pdf",width=15,height=10)
par(mfrow=c(2,3))
neff = matrix(0,n,3)
mx.sis = matrix(0,n,3)
for (i in 1:3){
  sis = SIS(y[,i],alpha,beta,sig2,tau2,m0,C0,M)
  qsis = sis$qs
  neff[,i] = sis$neff
  ts.plot(qsis,col=c(3,2,3),main=paste("sig2=",sig2s[i],sep=""),xlab="time")
  lines(qdlm[i,,1])
  lines(qdlm[i,,2])
  lines(qdlm[i,,3])
  mx.sis[,i] = sis$mx
}
for (i in 1:3){
  plot(neff[,i],ylim=c(0,M),xlab="time",ylab="Effective sample size")
}
dev.off()

pdf(file="rmae3.pdf",width=15,height=10)
par(mfrow=c(2,3))
neff = matrix(0,n,3)
mx.bf = matrix(0,n,3)
for (i in 1:3){
  bf = BF(y[,i],alpha,beta,sig2,tau2,m0,C0,M)
  qbf = bf$qs
  neff[,i] = bf$neff
  ts.plot(qbf,col=c(3,2,3),main=paste("sig2=",sig2s[i],sep=""),xlab="time")
  lines(qdlm[i,,1])
  lines(qdlm[i,,2])
  lines(qdlm[i,,3])
  mx.bf[,i] = bf$mx
}
for (i in 1:3){
  plot(neff[,i],ylim=c(0,M),xlab="time",ylab="Effective sample size")
}
dev.off()

pdf(file="rmae4.pdf",width=15,height=10) 
par(mfrow=c(2,3))
neff = matrix(0,n,3)
perc = 0.2
mx.sis1 = matrix(0,n,3)
for (i in 1:3){
  sis1 = SIS1(y[,i],alpha,beta,sig2,tau2,m0,C0,M,perc)
  qsis1 = sis1$qs
  neff[,i] = sis1$neff
  ts.plot(qsis1,col=c(3,2,3),main=paste("sig2=",sig2s[i],sep=""),xlab="time")
  lines(qdlm[i,,1])
  lines(qdlm[i,,2])
  lines(qdlm[i,,3])
  mx.sis1[,i] = sis1$mx
}
for (i in 1:3){
  plot(neff[,i],ylim=c(0,M),xlab="time",ylab="Effective sample size")
  abline(h=perc*M,col=2)
}
dev.off()

pdf(file="rmae5.pdf",width=15,height=5)
par(mfrow=c(1,3))
for (i in 1:3){
  ts.plot(cbind(mx.sis[,i],mx.bf[,i],mx.sis1[,i]),col=1:3)
  lines(qdlm[i,,2],col=4)
}
dev.off()


# Larger simulated exercise
# ------------------------
nrep  = 100
nrep  = 10
M     = 1000
n     = 500
sig2s = c(0.25,0.5,1.0)
alpha = 0.0
beta  = 1.0
tau2  = 0.2
x0    = 0.0
m0    = 0.0
C0    = 10
percs = c(0.2,0.5,0.8)
nperc = length(percs)

set.seed(12345)
rmae = array(0,c(nrep,3,3))
for (l in 1:nperc){
  perc = percs[l]
  for (r in 1:nrep){
    print(c(r,perc))
    for (i in 1:3){
      sig2  = sig2s[i]
      sig   = sqrt(sig2)
      tau   = sqrt(tau2)
      sim   = DLM.SIM(n,tau,sig,x0)
      y     = sim$y
      dlm   = DLM(y,alpha,beta,sig2,tau2,m0,C0)$m
      bf    = BF(y,alpha,beta,sig2,tau2,m0,C0,M)$qs[,2]
      sis   = SIS1(y,alpha,beta,sig2,tau2,m0,C0,M,perc)$qs[,2]
      rmae[r,l,i] = mean(abs(sis-dlm))/mean(abs(bf-dlm))
    }
  }
}

pdf(file="rmae-resampling-n500.pdf",width=12,height=10)
par(mfrow=c(3,nperc))
for (i in 1:3){
  for (l in 1:nperc){
    boxplot(rmae[,l,i],xlab=paste("sig2=",sig2s[i],sep=""),
            ylab="Relative MAE",outline=FALSE,ylim=range(rmae[,,i]))
    abline(h=1,col=2)
    title(paste("M1=",percs[l]*M,sep=""))
  }
}
dev.off()