###########################################################################
#
# The first order dynamic linear model (aka the local level model):
#
#      y(t)    = beta(t)   + v(t)      v(t) iid N(0,V)
#      beta(t) = beta(t-1) + w(t)      w(t) iid N(0,W)
#
# Posterior inference based on a few functions/packages.
#
# Routine mcmc.joint + ffbs are mine and perform full posterior inference 
# for beta(1),...,beta(n),V,W via MCMC throught the FFBS scheme.
#
# Additional R packages:  dlm (by Giovanni Petris) & bsts (by Steve Scott)
#
# We also use Silvaneo Santos's R package kDGLM for filtering.  
# This package also entertains dynamid generalized linear structures.
#
#
# Note:  I would like to thank Silvaneo dos Santos Jr., who maintains
#        the R package kDGLM, for helping me out organizing the final 
#        part of this script.
#
#
# Copyright: Hedibert F. Lopes, PhD
#            Professor of Statistics and Econometrics
#            Insper Institute of Education and Research
#  
# Date: March 16th 2026
#
###########################################################################

sample.data = function(n,W,V,beta0){
  sdW     = sqrt(W)
  sdV     = sqrt(V)
  beta    = rep(0,n)
  y       = rep(0,n)
  beta[1] = rnorm(1,beta0,sdW)
  y[1]    = rnorm(1,beta[1],sdV)
  for (t in 2:n){
    beta[t] = rnorm(1,beta[t-1],sdW)  
    y[t]    = rnorm(1,beta[t],sdV)
  }
  return(y)
}

mcmc.joint = function(y,a1,R1,av,bv,aw,bw,M0,M,V,W){
  n      = length(y)
  V.draw = V
  W.draw = W
  niter  = M0+M
  draws  = matrix(0,niter,n+2)
  for (iter in 1:niter){
    beta.draw = ffbs(y,rep(1,n),0.0,V.draw,1.0,0.0,W.draw,a1,R1,nd=1)
    par2v     = bv+sum((y-beta.draw)^2)/2
    par2w     = bw+sum((beta.draw[2:n]-beta.draw[1:(n-1)])^2)/2
    V.draw    = 1/rgamma(1,av+n/2,par2v)
    W.draw    = 1/rgamma(1,aw+(n-1)/2,par2w)
    draws[iter,] = c(beta.draw,V.draw,W.draw)
  }
  return(draws[(M0+1):niter,])
}

ffbs = function(y,Ft,alpha,V,G,gamma,W,a1,R1,nd=1){
  n = length(y)
  if (length(Ft)==1){Ft=rep(Ft,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]+Ft[1]*a[1]
  Q    = R[1]*Ft[1]**2+V[1]
  A    = R[1]*Ft[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]+Ft[t]*a[t]
    Q    = R[t]*Ft[t]**2+V[t]
    A    = R[t]*Ft[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
  }
}

set.seed(123456)
n     = 300
W     = 0.5
V     = 1.0
beta0 = 0
y     = sample.data(n,W,V,beta0)
y     = y[100:n]
n     = length(y)
ts.plot(y)
  


# Prior for beta(0) ~ N(m0,C0)
# Prior for beta(1) ~ N(a1,R1), a1=m0,R1=C0+W
a1     = 0
R1     = 10
av     = 2.01
bv     = 1.01
aw     = 2.01
bw     = 1.01*W    
m0     = 0
C0     = 10-W

# MCMC initual values
V0     = 1
W0     = 0.5
M0     = 10000
M      = 10000

set.seed(65432)
draws = mcmc.joint(y,a1,R1,av,bv,aw,bw,M0,M,V0,W0)
qbeta = t(apply(draws[,1:n],2,quantile,c(0.05,0.5,0.95)))

par(mfrow=c(1,2))
hist(draws[,n+1],main="V",prob=TRUE,xlab="")
hist(draws[,n+2],main="W",prob=TRUE,xlab="")

par(mfrow=c(1,1))
ts.plot(qbeta,col=c(3,2,3),lwd=2)
points(y,pch=16,cex=0.5)


############################################
# Giovanni Petris' R package dlm
############################################
fit.dlm = StructTS(y, "level") 

fit.dlm

tab = cbind(round(fit.dlm$coef,3),c(W,V))
colnames(tab) = c("Estimated","True")
rownames(tab) = c("W","V")
tab


par(mfrow=c(1,1))
ts.plot(qbeta[,2],lwd=2)
lines(fit.dlm$fitted,col=2)

############################################
# Steve Scott's R package bsts
############################################
library(bsts)
ss = AddLocalLinearTrend(list(), y) 
model = bsts(y, state.specification = ss, niter = 10000)

sigmas = cbind(model$sigma.obs, model$sigma.trend.level) 

tab = round(cbind(c(V,W),apply(sigmas[5001:10000,],2,mean), 
                  apply(sigmas[5001:10000,],2,median), 
                  apply(sigmas[5001:10000,],2,quantile,0.025), 
                  apply(sigmas[5001:10000,],2,quantile,0.975)),2)
rownames(tab) = c("V","W")
colnames(tab) = c("True","Mean","Median","2.5%","97.5")
tab

names = c("V","W")
par(mfrow=c(2,3)) 
for (i in 1:2){
  ts.plot(sigmas[,i],ylab="",main=names[i]) 
  abline(v=5000,col=2,lwd=2)
  acf(sigmas[5001:10000,i],main="") 
  hist(sigmas[5001:10000,i],xlab="",main="",prob=TRUE)
}

trend = model$state.contributions[5001:10000,1,]

qtrend = t(apply(trend,2,quantile,c(0.025,0.5,0.975)))

par(mfrow=c(1,1))
ts.plot(qbeta[,2],lwd=2,ylim=range(qbeta,fit.dlm$fitted,qtrend))
lines(qbeta[,1],lwd=2)
lines(qbeta[,3],lwd=2)
lines(fit.dlm$fitted,col=2,lwd=2)
lines(qtrend[,1],col=3,lwd=2)
lines(qtrend[,2],col=3,lwd=2)
lines(qtrend[,3],col=3,lwd=2)
legend("topleft",legend=c("mcmc.joint (we know all derivations!)",
                          "dlm (what prior for (V,W)?)",
                          "bsts (what prior for (V,W)?)"),
       col=1:3,lwd=2,bty="n")



############################################
# Silvaneo Santos's R package kDGLM
############################################
#remotes::install_github('silvaneojunior/kDGLM')
library("kDGLM")

data = data.frame(y=y)

# Unknown but constant observational variance
fit.kdglm =  kdglm(y~pol(D=0.7,name='nivel.media'),
                   V=~pol(D=1,name='nivel.var'), 
                   family=Normal,data=data)

fit.kdglm$var.names  # Names for the latent states (theta)
fit.kdglm$pred.names # Names for the linear predictions (lambda)

# A few plots
# For models with more than one outcome, the "outcome" argument defines 
# which outcomes should be presented (multiple values can be selected)

# Lag=-1 -> smoothed distributions, 95% C.I (default)
plot(fit.kdglm,outcome='y',lag=-1 , pred.cred = 0.95) 

# Lag=0 -> Filtered distributions
plot(fit.kdglm,lag=0) 

# Lag=k -> k steps ahead forecast (k>0)
plot(fit.kdglm,lag=1)

# Linear predictors
# The argument linear.predictors defines which linear predictors should be plotted (You can choose more than one!)
# 'y' returns the linear predictor for the mean, 'V' returns the linear predictor for the log-variance
plot(fit.kdglm,linear.predictors ='y',lag=-1) # smoothed distribution


# Latent states
# The argument latent.states defines which latent states should be plotted (You can choose more than one!)
# In the current example this is irrelevant as the latent state is also the linear predictor since 
# the design matrix F is the identity.
# For a more general model
# If mu = a + b + c, the argument latent.states would be used to evaluate
# the states a, b and c separately, while the argumento linear.predictor would be used
# to evaluate mu
plot(fit.kdglm,latent.states ='nivel.media',lag=-1) # smoothed distribution


# O get the mean vectors and covariance matrices, one case use the coef function
coefs=coef(fit.kdglm,lag=-1)

# Latent states
coefs$theta.mean
coefs$theta.cov

# Linear predictors
coefs$lambda.mean
coefs$lambda.cov

m = coefs$theta.mean[1,]
s = sqrt(coefs$theta.cov[1,1,])
L = m-2*s
U = m+2*s

par(mfrow=c(1,1),mar = c(2, 2, 1, 1))
plot(m,type="l",ylim=range(L,U,y),col=2)
lines(L,col=3)
lines(U,col=3)
points(y)