# Download the authors' R code to do filtering and smoothing
#source("http://www.stat.pitt.edu/stoffer/tsa2/Rcode/New6/ssall.R")
#or just
source("ssall.R")
# since this is in the directory of R scripts on the course website

# ssall is reprinted at the bottom of this script

# R code for this example is in 
# http://www.stat.pitt.edu/stoffer/tsa2/Rcode/New6/ex65.txt:

#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!#
#--   assumes ss0 has been sourced    --#
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!#


#-- generate data --
set.seed(1)  
num=50
w = rnorm(num+1,0,1)
v = rnorm(num,0,1)
mu = array0(cumsum(w), offset=0) # state: mu[0],mu[1],...,mu[50]    
y=mu[1:num]+v                    # obs:          y[1],..., y[50]

# set parameters
mu0=0
sigma0=1
phi=1
cQ=1   #  These are the Cholesky decomps of Q and R, (eg  Q=t(cQ)%*%cQ, where cQ is upper triangular)
cR=1   #  which are the standard devs in this case.

  # you don't have to run the filter because Ksmooth0 returns 
  # the filters and predictions as well as the smoothers...
  #  ... that's why the kf= line below is commented out ...
  #  it's just an example of what to do if you don't want to smooth
# kf = Kfilter0(50,y,1,mu0,sigma0,phi,cQ,cR) 
ks = Ksmooth0(50,y,1,mu0,sigma0,phi,cQ,cR)   

# To see the effect of NOT smoothing, replace ks = ... above by
#ks = Kfilter0(50,y,1,mu0,sigma0,phi,cQ,cR)
  
#-- pictures
##  notes:
## mu_t^t-1=ks$xp,  P_t^t-1=ks$Pp,   t=  1,...,n and "p" for prediction
## mu_t^t=ks$xf,    P_t^t=ks$Pf,     t=0,1,...,n and "f" for filter
## mu_t^n=ks$xs,    P_t^n=ks$Ps,     t=0,1,...,n and "s" for smoother
  Time=0:num
  par(mfrow=c(3,1))
  plot(Time[-1], mu[-1], main="Prediction", ylim=c(-5,10))  #[-1] leaves off first component
   lines(ks$xp)
   lines(ks$xp+2*sqrt(ks$Pp), lty="dashed", col="blue")
   lines(ks$xp-2*sqrt(ks$Pp), lty="dashed", col="blue")
  plot(Time, mu, main="Filter", ylim=c(-5,10))
   lines(Time, ks$xf)
   lines(Time, ks$xf+2*sqrt(ks$Pf), lty="dashed", col="blue")
   lines(Time, ks$xf-2*sqrt(ks$Pf), lty="dashed", col="blue")
  plot(Time, mu,  main="Smoother", ylim=c(-5,10))
   lines(Time, ks$xs)
   lines(Time, ks$xs+2*sqrt(ks$Ps), lty="dashed", col="blue")
   lines(Time, ks$xs-2*sqrt(ks$Ps), lty="dashed", col="blue")







###################################
#############################################################
#     This is ssall.R, where you get it all and more!
# 
#        You get: Kfilter0, Ksmooth0  (level 0 - see below)
#                 Kfilter1, Ksmooth1  (level 1 - see below)
#                 Kfilter2, Ksmooth2  (level 2 - see below)
#                 and  Oarray  is included       
#############################################################
#  LEVEL 0: state space modeling (easy stuff, no inputs)
#           the filter and smoother with fixed measurement 
#           matrix A(t)= A for all t
#
#  LEVEL 1: state space modeling 
#           (i) allows inputs, u(t) 
#           (ii) the measurement matrix is a sequence A(t)
#           --read the notes in the filter for special cases--
#
#      LEVELS 0 and 1 use the Oarray package (included here)
#
#  LEVEL 2: state space modeling  
#          (i) allows inputs, u(t) 
#          (ii) the measurement matrix is a sequence A(t)
#          (iii) errors can be correlated
#
# !!!! check out the notes for each function below for particulars  !!!!!
#
#############################################################
#############################################################
#********************* start LEVEL 0 ************************
#############################################################
#   the filter
#############################################################
Kfilter0 = function(num,y,A,mu0,Sigma0,Phi,cQ,cR){
  #
  # NOTE: must give cholesky decomp: cQ=chol(Q), cR=chol(R)
 Q=t(cQ)%*%cQ
 R=t(cR)%*%cR
   # y is num by q  (time=row series=col)
   # A is a q by p matrix
   # R is q by q
   # mu0 is p by 1
   # Sigma0, Phi, Q are p by p
 Phi=as.matrix(Phi)
 pdim=nrow(Phi)    
 y=as.matrix(y)
 qdim=ncol(y)
 xp=array(NA, dim=c(pdim,1,num))         # xp=x_t^{t-1}          
 Pp=array(NA, dim=c(pdim,pdim,num))      # Pp=P_t^{t-1}
 xf=array0(NA, dim=c(pdim,1,num+1), offset=c(1,1,0))      # xf=x_t^t
 Pf=array0(NA, dim=c(pdim,pdim,num+1), offset=c(1,1,0))   # Pf=P_t^t
 innov=array(NA, dim=c(qdim,1,num))      # innovations
 sig=array(NA, dim=c(qdim,qdim,num))     # innov var-cov matrix
 like=0                                  # -log(likelihood)
 xf[,,0]=mu0
 Pf[,,0]=Sigma0
 for (i in 1:num){
  xp[,,i]=Phi%*%xf[,,i-1]
  Pp[,,i]=Phi%*%Pf[,,i-1]%*%t(Phi)+Q
    siginv=A%*%Pp[,,i]%*%t(A)+R
  sig[,,i]=(t(siginv)+siginv)/2     # make sure sig is symmetric
    siginv=solve(sig[,,i])          # now siginv is sig[[i]]^{-1}
  K=Pp[,,i]%*%t(A)%*%siginv
  innov[,,i]=y[i,]-A%*%xp[,,i]
  xf[,,i]=xp[,,i]+K%*%innov[,,i]
  Pf[,,i]=Pp[,,i]-K%*%A%*%Pp[,,i]
    sigmat=as.matrix(sig[,,i], nrow=qdim, ncol=qdim)
  like= like + log(det(sigmat)) + t(innov[,,i])%*%siginv%*%innov[,,i]
  }
    like=0.5*like
    list(xp=xp,Pp=Pp,xf=xf,Pf=Pf,like=like,innov=innov,sig=sig,Kn=K)
}
#
##########################################################
#  the smoother
###########################################################
Ksmooth0 = function(num,y,A,mu0,Sigma0,Phi,cQ,cR){
#
# Note: Q and R are given as Cholesky decomps
#       cQ=chol(Q), cR=chol(R)
#
 kf=Kfilter0(num,y,A,mu0,Sigma0,Phi,cQ,cR)
 pdim=nrow(as.matrix(Phi))  
 xs=array0(NA, dim=c(pdim,1,num+1), offset=c(1,1,0))      # xs=x_t^n
 Ps=array0(NA, dim=c(pdim,pdim,num+1), offset=c(1,1,0))   # Ps=P_t^n
 J=array0(NA, dim=c(pdim,pdim,num), offset=c(1,1,0))      # J=J_t
 xs[,,num]=kf$xf[,,num] 
 Ps[,,num]=kf$Pf[,,num]
 for(k in num:1)  {
 J[,,k-1]=(kf$Pf[,,k-1]%*%t(Phi))%*%solve(kf$Pp[,,k])
 xs[,,k-1]=kf$xf[,,k-1]+J[,,k-1]%*%(xs[,,k]-kf$xp[,,k])
 Ps[,,k-1]=kf$Pf[,,k-1]+J[,,k-1]%*%(Ps[,,k]-kf$Pp[,,k])%*%t(J[,,k-1])
}
list(xs=xs,Ps=Ps,J=J,xp=kf$xp,Pp=kf$Pp,xf=kf$xf,Pf=kf$Pf,like=kf$like,Kn=kf$K)
} 
################## end level 0 ##############################
#************************************************************
#******************* start LEVEL 1 **************************
#############################################################
#   the filter
#############################################################
Kfilter1 = function(num,y,A,mu0,Sigma0,Phi,Ups,Gam,cQ,cR,input){
  #
  # NOTE: must give cholesky decomp: cQ=chol(Q), cR=chol(R)
 Q=t(cQ)%*%cQ
 R=t(cR)%*%cR
   # y is num by q  (time=row series=col)
   # input is num by r (use 0 if not needed)
   # A is an array with dim=c(q,p,num)
   # Ups is p by r  (use 0 if not needed)
   # Gam is q by r  (use 0 if not needed)
   # R is q by q
   # mu0 is p by 1
   # Sigma0, Phi, Q are p by p
 Phi=as.matrix(Phi)
 pdim=nrow(Phi)     
 y=as.matrix(y)
 qdim=ncol(y)
 rdim=ncol(as.matrix(input))
  if (Ups==0) Ups = matrix(0,pdim,rdim)
  if (Gam==0) Gam = matrix(0,qdim,rdim)
 ut=matrix(input,num,rdim)
 xp=array(NA, dim=c(pdim,1,num))      # xp=x_t^{t-1}          
 Pp=array(NA, dim=c(pdim,pdim,num))   # Pp=P_t^{t-1}
 xf=array0(NA, dim=c(pdim,1,num+1), offset=c(1,1,0))       # xf=x_t^t
 Pf=array0(NA, dim=c(pdim,pdim,num+1), offset=c(1,1,0))    # Pf=x_t^t
 innov=array(NA, dim=c(qdim,1,num))   # innovations
 sig=array(NA, dim=c(qdim,qdim,num))  # innov var-cov matrix
 like=0                               # -log(likelihood)
 xf[,,0]=mu0
 Pf[,,0]=Sigma0
 for (i in 1:num){
  xp[,,i]=Phi%*%xf[,,i-1] + Ups%*%ut[i,]
  Pp[,,i]=Phi%*%Pf[,,i-1]%*%t(Phi)+Q
      B = matrix(A[,,i], nrow=qdim, ncol=pdim)  
      siginv=B%*%Pp[,,i]%*%t(B)+R
  sig[,,i]=(t(siginv)+siginv)/2     # make sure sig is symmetric
    siginv=solve(sig[,,i])          # now siginv is sig[[i]]^{-1}
  K=Pp[,,i]%*%t(B)%*%siginv
  innov[,,i]=y[i,]-B%*%xp[,,i]-Gam%*%ut[i,]
  xf[,,i]=xp[,,i]+K%*%innov[,,i]
  Pf[,,i]=Pp[,,i]-K%*%B%*%Pp[,,i]
    sigmat=matrix(sig[,,i], nrow=qdim, ncol=qdim)
  like= like + log(det(sigmat)) + t(innov[,,i])%*%siginv%*%innov[,,i]
  }
    like=0.5*like
    list(xp=xp,Pp=Pp,xf=xf,Pf=Pf,like=like,innov=innov,sig=sig,Kn=K)
}
#
##########################################################
#  the smoother
###########################################################
Ksmooth1 = function(num,y,A,mu0,Sigma0,Phi,Ups,Gam,cQ,cR,input){
#
# Note: Q and R are given as Cholesky decomps
#       cQ=chol(Q), cR=chol(R)
#  Use Ups=0 or Gam=0 or input=0 if these aren't needed
#
 kf=Kfilter1(num,y,A,mu0,Sigma0,Phi,Ups,Gam,cQ,cR,input)
 pdim=nrow(as.matrix(Phi))  
 xs=array0(NA, dim=c(pdim,1,num+1), offset=c(1,1,0))      # xs=x_t^n
 Ps=array0(NA, dim=c(pdim,pdim,num+1), offset=c(1,1,0))      # Ps=P_t^n
 J=array0(NA, dim=c(pdim,pdim,num), offset=c(1,1,0))      # J=J_t
 xs[,,num]=kf$xf[,,num] 
 Ps[,,num]=kf$Pf[,,num]
 for(k in num:1)  {
 J[,,k-1]=(kf$Pf[,,k-1]%*%t(Phi))%*%solve(kf$Pp[,,k])
 xs[,,k-1]=kf$xf[,,k-1]+J[,,k-1]%*%(xs[,,k]-kf$xp[,,k])
 Ps[,,k-1]=kf$Pf[,,k-1]+J[,,k-1]%*%(Ps[,,k]-kf$Pp[,,k])%*%t(J[,,k-1])
}
list(xs=xs,Ps=Ps,J=J,xp=kf$xp,Pp=kf$Pp,xf=kf$xf,Pf=kf$Pf,like=kf$like,Kn=kf$K)
} 
################## end level 1 ###################################
#*****************************************************************
#******************** start LEVEL 2 ******************************
##################################################################
#   the filter
##################################################################
Kfilter2 = function(num,y,A,mu1,Sigma1,Phi,Ups,Gam,cQ,cR,S,input){
  #
  ######## Reference Property 6.5 in Section 6.6 ###########
  # num is the number of observations
  # y is the data matrix (num by q)
  # A has to be an array with dim=c(q,p,num)
  # mu1= E(x_1) = x_1^0 = Phi%*%mu0 + Gam%*%input0 
  # Sigma1 = var(x_1)= P_1^0 =  Phi%*%Sigma0%*%t(Phi)+Q
  # input has to be a matrix (num by r) - similar to obs y
  # set Ups or Gam or input to 0 if not used
  # Must give Cholesky decomp: cQ=chol(Q), cR=chol(R)  
Q=t(cQ)%*%cQ
R=t(cR)%*%cR 
  #
 Phi=as.matrix(Phi)
 pdim=nrow(Phi) 
 y=as.matrix(y)
 qdim=ncol(y)
 rdim=ncol(as.matrix(input))
  if (Ups==0) Ups = matrix(0,pdim,rdim)
  if (Gam==0) Gam = matrix(0,qdim,rdim)
 ut=matrix(input,num,rdim)
 xp=array(NA, dim=c(pdim,1,num+1))      # xp=x_t^{t-1}          
 Pp=array(NA, dim=c(pdim,pdim,num+1))   # Pp=P_t^{t-1}
 xf=array(NA, dim=c(pdim,1,num+1))      # xf=x_t^{t} 
 Pf=array(NA, dim=c(pdim,pdim,num+1))   # Pf=P_t^{t}
 Gain=array(NA, dim=c(pdim,qdim,num))
 innov=array(NA, dim=c(qdim,1,num))   # innovations
 sig=array(NA, dim=c(qdim,qdim,num))  # innov var-cov matrix
 like=0                               # -log(likelihood)
 xp[,,1]=mu1
 Pp[,,1]=Sigma1
 for (i in 1:num){
	B = matrix(A[,,i], nrow=qdim, ncol=pdim) 
   innov[,,i] = y[i,]-B%*%xp[,,i]-Gam%*%ut[i,]
    sigma = B%*%Pp[,,i]%*%t(B)+R 
    sigma=(t(sigma)+sigma)/2     # make sure sig is symmetric
   sig[,,i]=sigma
    siginv=solve(sigma)
   Gain[,,i]=(Phi%*%Pp[,,i]%*%t(B)+S)%*%siginv 
    K=Gain[,,i]
   xp[,,i+1]=Phi%*%xp[,,i] + Ups%*%ut[i,] + K%*%innov[,,i]
   Pp[,,i+1]=Phi%*%Pp[,,i]%*%t(Phi)+ Q - K%*%sig[,,i]%*%t(K)
   xf[,,i]=xp[,,i]+ Pp[,,i]%*%t(B)%*%siginv%*%innov[,,i]
   Pf[,,i]=Pp[,,i] - Pp[,,i]%*%t(B)%*%siginv%*%B%*%Pp[,,i] 
     sigma=matrix(sigma, nrow=qdim, ncol=qdim)
   like= like + log(det(sigma)) + t(innov[,,i])%*%siginv%*%innov[,,i]
  }	 
  like=0.5*like
  list(xp=xp,Pp=Pp,xf=xf,Pf=Pf, K=Gain,like=like,innov=innov,sig=sig)
}
#
##########################################################
#  the smoother
###########################################################
Ksmooth2 = function(num,y,A,mu1,Sigma1,Phi,Ups,Gam,cQ,cR,S,input){
#
# Note: Q and R are given as Cholesky decomps
#       cQ=chol(Q), cR=chol(R)
#  Use Ups=0 or Gam=0 or input=0 if these aren't needed
#
 kf=Kfilter2(num,y,A,mu1,Sigma1,Phi,Ups,Gam,cQ,cR,S,input)
 pdim=nrow(as.matrix(Phi))  
 xs=array(NA, dim=c(pdim,1,num))      # xs=x_t^n
 Ps=array(NA, dim=c(pdim,pdim,num))   # Ps=P_t^n
 J=array(NA, dim=c(pdim,pdim,num))    # J=J_t
 xs[,,num]=kf$xf[,,num]      
 Ps[,,num]=kf$Pf[,,num]
 for(k in num:2)  {
 J[,,k-1]=(kf$Pf[,,k-1]%*%t(Phi))%*%solve(kf$Pp[,,k])
 xs[,,k-1]=kf$xf[,,k-1]+J[,,k-1]%*%(xs[,,k]-kf$xp[,,k])
 Ps[,,k-1]=kf$Pf[,,k-1]+J[,,k-1]%*%(Ps[,,k]-kf$Pp[,,k])%*%t(J[,,k-1])
}
list(xs=xs,Ps=Ps,J=J,xp=kf$xp,Pp=kf$Pp,xf=kf$xf,Pf=kf$Pf,like=kf$like,Kn=kf$K)
} 
################## end ss2 #######################################
#*****************************************************************
#\\\\\\\\\\//////////////////\\\\\\\\\\\\\\\\\////////////////\\\\ 
##################################################################
### Oarray stuff so you don't have to load it each time  #########
##################################################################
"Oarray" <-
function(data=NA, dim=length(data), dimnames=NULL, offset=rep(1, length(dim)),
  drop.negative=TRUE)
{
  if (!is.numeric(offset) || length(offset) != length(dim))
    stop("\"offset\" must be numeric vector with same length as \"dim\"")
  if (drop.negative && any(offset < 0))
    stop("Non-negative offsets only")

  robj <- array(data = data, dim = dim, dimnames = dimnames)
  attr(robj, "offset") <- offset
  attr(robj, "drop.negative") <- drop.negative
  class(robj) <- "Oarray"
  robj
}

"as.Oarray" <- function(x, offset=rep(1, length(dim)), drop.negative=TRUE)
{
  x <- as.array(x)
  dim <- dim(x)
  Oarray(x, dim = dim, dimnames = dimnames(x), offset = offset,
    drop.negative = drop.negative)
}

"as.array.default" <- get("as.array", pos=grep("package:base",
  search()), mode="function")

"as.array" <- function(x)
  UseMethod("as.array")

"as.array.Oarray" <- function(x)
{
  x <- unclass(x)
  attr(x, "offset") <- NULL
  attr(x, "drop.negative") <- NULL
  NextMethod(x)
}

"is.Oarray" <- function(x)
  inherits(x, "Oarray") && !is.null(attr(x, "offset")) &&
    !is.null(attr(x, "drop.negative"))

# this function takes numeric index sets from the original call
# and maps them using the offset: note that drop=FALSE only works
# if provided as the final argument

".handleTheOffset" <- function(mc, dim, offset, dn)
{
  for (i in seq(along=dim)) {
    ii <- mc[[2+i]]

    if (missing(ii)) next

    if (is.symbol(ii) || is.call(ii))
      ii <- eval.parent(ii, 3)

    if (is.numeric(ii)) {

      if (!dn || all(ii>=0))
        ii <- ifelse(ii>=offset[i], ii - offset[i] + 1, dim[i]+1)
      else {
        if (all(ii <= -offset[i]))
          ii <- ii + offset[i] - 1
        else stop("subscript out of bounds")
      }

      mc[[2+i]] <- ii
    }
  }
  mc
}

"[.Oarray" <- function(x, ...)
{
  mc <- match.call()
  k <- length(mc)
  offset <- attr(x, "offset")
  dn <- attr(x, "drop.negative")
  dim <- dim(x)

  if (k==3 && mc[[3]]=="")
    return(as.array(x))

  if (k < 2+length(dim))
    stop("incorrect number of dimensions")

  mc <- .handleTheOffset(mc, dim, offset, dn)
  mc[[1]] <- as.name("[")
  mc[[2]] <- as.name("x")
  x <- as.array(x)
  eval(mc)
}

"[<-.Oarray" <- function(x, ..., value)
{
  mc <- match.call()
  k <- length(mc)
  offset <- attr(x, "offset")
  dn <- attr(x, "drop.negative")
  dim <- dim(x)

  if (k==4 && mc[[3]]=="")
    return(Oarray(value, dim, dimnames(x), offset, dn))

  if (k < 3+length(dim))
    stop("incorrect number of dimensions")

  mc <- .handleTheOffset(mc, dim, offset, dn)
  mc[[1]] <- as.name("[<-")
  mc[[2]] <- as.name("x")
  x <- as.array(x)
  robj <- eval(mc)
  Oarray(robj, dim, dimnames(x), offset, dn)
}

"print.Oarray" <-
function(x, ...)
{
  d <- dim(x)
  dn <- dimnames(x)
  if (is.null(dn))
    dn <- vector("list", length(d))
  offset <- attr(x, "offset")
  x <- as.array(x)

  for (i in seq(along=dn))
    if (is.null(dn[[i]])) {
      dn[[i]] <- 0:(d[i]-1) + offset[i]
      if (i==1 || i==2) {
        dn[[i]] <- format(dn[[i]])
        if (i==1)
          dn[[i]] <- paste("[", dn[[i]], ",]", sep="")
        else
          dn[[i]] <- paste("[,", dn[[i]], "]", sep="")
      }
    }
  dimnames(x) <- dn
  NextMethod("print")
}
array0<-Oarray
cat("  You now have", "\n")
cat("  Kfilter0, Ksmooth0", "\n")
cat("  Kfilter1, Ksmooth1", "\n")
cat("  Kfilter2, Ksmooth2", "\n")
cat("  and Oarray", "\n")
# end ###################################