######################################################
# In our 3rd lab session we have implemented both
# Kalman filter and Kalman smoother for the simple
# first order dynamic linear model, also known as
# the local level model.
######################################################

kalmanfilter = function(y,m0,C0,sig2,tau2){
	n=length(y)
	R=rep(0,n);Q=rep(0,n);A=rep(0,n)
	m=rep(0,n);C=rep(0,n);B=rep(0,n)
	R[1] = C0 + tau2
	Q[1] = R[1] + sig2
	A[1] = R[1]/Q[1]
	C[1] = A[1]*sig2
	m[1] = (1-A[1])*m0 + A[1]*y[1]
	B[1] = C[1]/(C[1]+tau2)
	for (t in 2:n){
		R[t] = C[t-1] + tau2
		Q[t] = R[t] + sig2
		A[t] = R[t]/Q[t]
		C[t] = A[t]*sig2
		m[t] = (1-A[t])*m[t-1] + A[t]*y[t]
		B[t] = C[t]/(C[t]+tau2)
	}
	return(list(m=m,C=C,B=B))
}

kalmansmoother = function(y,m,C,B){
	n = length(y)
	m1 = rep(0,n)
	C1 = rep(0,n)
	m1[n] = m[n]
	C1[n] = C[n]
	for (t in (n-1):1){
		m1[t] = (1-B[t])*m[t]+B[t]*m1[t+1]
		C1[t] = (1-B[t])*C[t]+B[t]^2*C1[t+1]
	}
	return(list(m1=m1,C1=C1))
}

marginallikelihood = function(y,m0,C0,sig2,tau2){
	n = length(y)
	m = m0
	C = C0
	like = 0.0	
	for (t in 1:n){
		R = C + tau2
		Q = R + sig2
		like = like + dnorm(y[t],m,sqrt(Q),log=TRUE)
		A = R/Q
		m = (1-A)*m + A*y[t]
		C = A*sig2
	}
	return(like)
}

#######################################################
# Simulated data
#######################################################
set.seed(123456)
n    = 50
sig2 = 1
tau2 = 0.5
x0   = 0
x    = rep(0,n)
y    = rep(0,n)
x[1] = rnorm(1,x0,sqrt(tau2))
y[1] = rnorm(1,x[1],sqrt(sig2))
for (t in 2:n){
  x[t] = rnorm(1,x[t-1],sqrt(tau2))
  y[t] = rnorm(1,x[t],sqrt(sig2))
  # Uncomment the following if you want a break in the states.
  #if (t==30){x[t]=20}
}
# Uncomment the following line if you want 
# an observational outlier.
#y[30]=-100

plot(y,pch=16,xlab="Time",ylab="")
lines(x,col=2,type="b",pch=16)
title("y (black) versus x (red)")

# Prior for the initial state
m0 =   0
C0 = 100

# Running Kalman filter and Kalman smoother
myrun = kalmanfilter(y,m0,C0,sig2,tau2)
m = myrun$m
C = myrun$C
B = myrun$B

myotherrun = kalmansmoother(y,m,C,B)
m1 = myotherrun$m1
C1 = myotherrun$C1

ts.plot(cbind(m-2*sqrt(C),m,m+2*sqrt(C)),lwd=2,type="b",ylab="State")
lines(m1,col=grey(0.5),lwd=2,type="b",pch=16)
lines(m1-2*sqrt(C1),col=grey(0.5),lwd=2,type="b",pch=16)
lines(m1+2*sqrt(C1),col=grey(0.5),lwd=2,type="b",pch=16)
title("Kalman filter (black) vs Kalman smoother (grey)")


# Marginal likelihood 
N  = 100
s2 = seq(0.1,2,length=N)
t2 = seq(0.1,2,length=N)
loglike = matrix(0,N,N)
for (i in 1:N)
  for (j in 1:N)
    loglike[i,j] = marginallikelihood(y,m0,C0,s2[i],t2[j])
like = exp(loglike-max(loglike))

sig2mle = s2[apply(like==max(like),1,sum)==1]
tau2mle = t2[apply(like==max(like),2,sum)==1]

contour(s2,t2,like,xlab=expression(sigma^2),ylab=expression(tau^2), 
        drawlabels=FALSE)
points(sig2,tau2,col=2,pch=16,cex=2)
title(paste("sig2 mle=",round(sig2mle,3)," -- tau2 mle=",round(tau2mle,3),sep=""))




#######################################################
# 1st order DLM for South Africa CPI
#######################################################
data = read.table("CPI-SA.txt",header=FALSE)
n    = nrow(data)
data = data[n:1,]
y    = data[,2]
ind  = seq(1,n,length=5)
dat  = data[ind,1]

par(mfrow=c(1,1))
plot(y,xlab="Months",ylab="CPI",axes=FALSE,type="l")
title("South African CPI")
axis(2);axis(1,at=ind,label=dat)

# Prior for the initial state
m0 =    10
C0 = 10000

# Marginal likelihood 
N  = 100
s2 = seq(0.01,0.1,length=N)
t2 = seq(0.01,1.0,length=N)
loglike = matrix(0,N,N)
for (i in 1:N)
  for (j in 1:N)
    loglike[i,j] = marginallikelihood(y,m0,C0,s2[i],t2[j])
like = exp(loglike-max(loglike))

sig2mle = s2[apply(like==max(like),1,sum)==1]
tau2mle = t2[apply(like==max(like),2,sum)==1]

contour(s2,t2,like,xlab=expression(sigma^2),ylab=expression(tau^2), 
        drawlabels=FALSE)
title(paste("sig2 mle=",round(sig2mle,3)," -- tau2 mle=",round(tau2mle,3),sep=""))


# Kalman filter and Kalman smoother
cpi.ff = kalmanfilter(data[,2],sig2mle,tau2mle,15,5)
m = cpi.ff$m
C = cpi.ff$C
B = cpi.ff$B
cpi.bs = kalmansmoother(data[,2],m,C,B)
m1 = cpi.bs$m1
C1 = cpi.bs$C1

par(mfrow=c(1,1))
L  = m-2*sqrt(C)
U  = m+2*sqrt(C)
L1 = m1-2*sqrt(C1)
U1 = m1+2*sqrt(C1)
plot(data[,2],xlab="Months",ylab="CPI",axes=FALSE,ylim=range(L,L1,U,U1),pch=16)
axis(2);axis(1,at=ind,label=dat)
lines(m,col=2)
lines(L,col=2)
lines(U,col=2)
lines(m1,col=4)
lines(L1,col=4)
lines(U1,col=4)