###################################################
# --- MUST source ss0 and EM0 and install nlme ---#
#source("http://www.stat.pitt.edu/stoffer/tsa2/Rcode/New6/ssall.R")
#source("http://www.stat.pitt.edu/stoffer/tsa2/Rcode/New6/EM0.R")

# or just:
source("ssall.R")
source("EM0.R")
# since these are in the directory of R scripts on the course website

# EMO is reprinted at the end

###################################################
#########################################
library(nlme)   # loads package nlme 
#########################################
#-- generate data (same as ex66, see that file for more details)
 set.seed(999)
 num=100
 N=num+1
 x = arima.sim(n=N, list(ar = .8, sd=1))  
 x = array0(x, dim=N, offset=0)
 v = rnorm(num,0,1)         
 y=ts(x[1:num]+v)              

#-- initial estimates 
 u=ts.intersect(y,lag(y,-1),lag(y,-2));
 varu=var(u);
 coru=cor(u);
 phi=coru[1,3]/coru[1,2];      # See ex6.3: phi=rho(2)/rho(1)         
 q = (1-phi^2)*varu[1,2]/phi;  # Ditto; obtained from ACF of y
 r = varu[1,1] - q/(1-phi^2);  # Ditto
 cr=sqrt(r)
 cq=sqrt(q)
 mu0=0
# Sigma0 = q/(1-phi^2)
 Sigma0 = 2.8 # Used by S&S; = 1/(1-(.8)^2); makes almost no difference in the end

#-- estimation (stop at 75 iterations or when -loglike changes by < .001)
#---  EM0 will display the iteration number, but you have to scroll manually to see it

em=EM0(num,y,1,mu0,Sigma0,phi,cq,cr,75,.001)    

# -- Standard Errors -- uses package nlme
phi=em$Phi
cq=chol(em$Q)
cr=chol(em$R)
mu0=em$mu0
Sigma0=em$Sigma0
para=c(phi,cq,cr)

#-- evaluates likelihood at estimates
 Linn=function(para){
  kf = Kfilter0(num,y,1,mu0,Sigma0,para[1],para[2],para[3])
  return(kf$like)
  }
emhess=fdHess(para, function(para) Linn(para))
stderr=sqrt(diag(solve(emhess$Hessian)))  
#-- display summary of estimation 
 estimate=c(para, em$mu0, em$Sigma0)
 se=c(stderr,NA,NA)
 u=cbind(estimate,se)
 rownames(u)=c("phi","sigw","sigv","mu0","Sigma0")
 u      
 





###########################################################
###########################################################
#  EM0 estimation with no missing obs, constant A matrix
#                and no inputs
#
#      !!!-- must source ss0 first --!!!
###########################################################
EM0 = function(num,y,A,mu0,Sigma0,Phi,cQ,cR,max.iter,tol){
#
# Note: Q and R are given as Cholesky decomps
#       cQ=chol(Q), cR=chol(R)
#
   Phi=as.matrix(Phi)
   pdim=nrow(Phi)
   y=as.matrix(y)
   qdim=ncol(y)
   cvg=1+tol
   like=matrix(0,max.iter,1)
   cat("iteration","   -loglikelihood", "\n")
#----------------- start EM -------------------------
for(iter in 1:max.iter){ 
  ks=Ksmooth0(num,y,A,mu0,Sigma0,Phi,cQ,cR)
  like[iter]=ks$like
  cat("   ",iter, "        ", ks$like, "\n")
  if(iter>1) cvg=like[iter-1]-like[iter]
  if(cvg<tol) break
  if(cvg<0) stop("Likelihood Not Increasing")
# Lag-One Covariance Smoothers 
  Pcs=array(NA, dim=c(pdim,pdim,num))     # Pcs=P_{t,t-1}^n
  eye=diag(1,pdim)
  Pcs[,,num]=(eye-ks$Kn%*%A)%*%Phi%*%ks$Pf[,,num-1]
   for(k in num:2){
   Pcs[,,k-1]=ks$Pf[,,k-1]%*%t(ks$J[,,k-2])+
           ks$J[,,k-1]%*%(Pcs[,,k]-Phi%*%ks$Pf[,,k-1])%*%t(ks$J[,,k-2])}
# Estimation
  S11=matrix(0,pdim,pdim)
  S10=matrix(0,pdim,pdim)
  S00=matrix(0,pdim,pdim)
  R=matrix(0,qdim,qdim)
  for(i in 1:num){
    S11 = S11 + ks$xs[,,i]%*%t(ks$xs[,,i]) + ks$Ps[,,i]
    S10 = S10 + ks$xs[,,i]%*%t(ks$xs[,,i-1]) + Pcs[,,i]
    S00 = S00 + ks$xs[,,i-1]%*%t(ks$xs[,,i-1]) + ks$Ps[,,i-1]
      u = y[i,]-A%*%ks$xs[,,i]
    R = R + u%*%t(u) + A%*%ks$Ps[,,i]%*%t(A)
  }
  Phi=S10%*%solve(S00)
  Q=(S11-Phi%*%t(S10))/num
   Q=(t(Q)+Q)/2        # make sure symmetric
  cQ=chol(Q)
  R=R/num
  cR=chol(R)
  mu0=ks$xs[,,0]
 # mu0=mu0              # uncomment this line to keep mu0 fixed
  Sigma0=ks$Ps[,,0]
 # Sigma0=Sigma0        # uncomment this line to keep Sigma0 fixed
}
 list(Phi=Phi,Q=Q,R=R,mu0=mu0,Sigma0=Sigma0,like=like[1:iter],niter=iter,cvg=cvg)
}