############################################################################################
# 
# AR(1) PLUS NOISE MODEL - SUBROUTINES
#
#      Observation equation: y(t)|x(t),theta   ~ N(x(t),sig2)
#      Evolution equation:   x(t)|x(t-1),theta ~ N(alpha+beta*x(t-1),tau2)
#  
# where x(0) ~ N(m0,C0) and theta=(alpha,beta,sig2,tau2) is known.
#
############################################################################################
#
# 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@ChicagoBooth.edu
# URL: http://faculty.chicagobooth.edu/hedibert.lopes/research/
#
############################################################################################

DLM.SIM = function(n,theta,x0){
  alpha = theta[1]
  beta  = theta[2]
  sig2  = theta[3]
  tau2  = theta[4]
  sig   = sqrt(sig2)
  tau   = sqrt(tau2)
  x     = rep(0,n)
  y     = rep(0,n)
  x[1]  = rnorm(1,alpha+beta*x0,tau)
  for (t in 2:n)
    x[t] = rnorm(1,alpha+beta*x[t-1],tau)
  y = rnorm(n,x,sig)  
  return(list(x=x,y=y))
}

DLM = function(y,theta,m0,C0){
  alpha = theta[1]
  beta  = theta[2]
  sig2  = theta[3]
  tau2  = theta[4]
  n     = length(y)
  m     = m0
  C     = C0
  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
    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))
} 

BF = function(y,theta,m0,C0,M){
  alpha = theta[1]
  beta  = theta[2]
  sig2  = theta[3]
  tau2  = theta[4]
  sig   = sqrt(sig2)
  tau   = sqrt(tau2)
  n     = length(y)
  x     = rnorm(M,m0,sqrt(C0))
  xs    = matrix(0,n,M)
  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=TRUE)
    x      = x[k]
    xs[t,] = x 
  }
  return(xs)
}

APF = function(y,theta,m0,C0,M){
  alpha = theta[1]
  beta  = theta[2]
  sig2  = theta[3]
  tau2  = theta[4]
  sig   = sqrt(sig2)
  tau   = sqrt(tau2)
  n     = length(y)
  x     = rnorm(M,m0,sqrt(C0))
  xs    = matrix(0,n,M)
  for (t in 1:n){
  	  a      = alpha+beta*x
  	  w      = dnorm(y[t],a,sig)
    k      = sample(1:M,size=M,prob=w,replace=TRUE)
  	  x      = rnorm(M,a[k],tau)
  	  w1     = dnorm(y[t],x,sig)/w[k]
  	  k      = sample(1:M,size=M,prob=w1,replace=TRUE)
    x      = x[k]
    xs[t,] = x 
  }
  return(xs)
}

OBF = function(y,theta,m0,C0,M){
  alpha = theta[1]
  beta  = theta[2]
  sig2  = theta[3]
  tau2  = theta[4]
  sig   = sqrt(sig2)
  tau   = sqrt(tau2)
  n  = length(y)
  w2 = sig2+tau2
  A  = tau2/w2
  s  = sqrt(w2)
  d  = sqrt(A*sig2)
  x  = rnorm(M,m0,sqrt(C0))
  xs = matrix(0,n,M)
  for (t in 1:n){
  	  a      = alpha+beta*x
    m      = (1-A)*a+A*y[t]
    x1     = rnorm(M,m,d)
    w      = dnorm(y[t],a,s)
    k      = sample(1:M,size=M,prob=w,replace=TRUE)
    x      = x1[k]
    xs[t,] = x 
  }
  return(xs)
}

OAPF = function(y,theta,m0,C0,M){
  alpha = theta[1]
  beta  = theta[2]
  sig2  = theta[3]
  tau2  = theta[4]
  sig   = sqrt(sig2)
  tau   = sqrt(tau2)
  n     = length(y)
  w2    = sig2+tau2
  A     = tau2/w2
  s     = sqrt(w2)
  d     = sqrt(A*sig2)
  x     = rnorm(M,m0,sqrt(C0))
  xs    = matrix(0,n,M)
  for (t in 1:n){
  	  a = alpha+beta*x
    w = dnorm(y[t],a,s)
    k = sample(1:M,size=M,prob=w,replace=TRUE)
    m = (1-A)*a[k]+A*y[t]
    x = rnorm(M,m,d)
    xs[t,] = x 
  }
  return(xs)
}