#logit and inverse logit transform
logit = function(x){log(x/(1-x))}
invlogit = function(x){exp(x)/(1+exp(x))}

set.seed(35345)
#Simulate data
alpha = 0.9;sigma=0.5;beta=1;tau=1
n = 100
y = rep(NA,n)
y[1] = rnorm(1,0,sigma/sqrt(1-alpha^2))
for(i in 2:n)
  y[i] = alpha*y[i-1]+rnorm(1,0,sigma)
z = beta*y+rnorm(n,0,tau)
#Inserting missing values
nNA = n/2
ind = sample(1:n,nNA)
z[ind] = NA

#Function for performing Kalman filtering
kalman = function(theta,z)
  {
    #theta is a vector of the parameters to be estimated
    #the standard deviations are log-transformed in order to get parameters on the real line
    #alpha is assumed to be in [0,1] and is transformed by the logit-transform
    alpha=invlogit(theta[1])
    sigma=exp(theta[2]);tau=exp(theta[3])
    sigma1=sigma/sqrt(1-alpha^2)
    n = length(z)
    z.hat = rep(NA,n);S=rep(NA,n)
    y.upd = rep(NA,n);P.upd = rep(NA,n)
    y.pred = rep(NA,n+1);P.pred = rep(NA,n+1)

    y.pred[1]=0;P.pred[1]=sigma1^2
    for(i in 1:n)
      {
        z.hat[i] = y.pred[i];S[i]=P.pred[i]+tau^2
        K = P.pred[i]/S[i]
        y.upd[i] = y.pred[i]+K*(z[i]-y.pred[i]);P.upd[i]=P.pred[i]*(1-K)
        if(is.na(z[i]))
           {
             y.upd[i] = y.pred[i];P.upd[i]=P.pred[i]
           }
        y.pred[i+1]=alpha*y.upd[i];P.pred[i+1]=alpha^2*P.upd[i]+sigma^2
      }
    l = sum(dnorm(z,z.hat,sqrt(S),log=TRUE),na.rm=TRUE)
    print(c(alpha,sigma,tau,l))
    list(y.upd=y.upd,P.upd=P.upd,loglik=l)
  }
#Calling kalman routine for true parameter sets and plotting results
theta = c(logit(alpha),log(sigma),log(tau))
foo = kalman(theta=theta,z=z)
matplot(cbind(1:n,1:n,1:n),cbind(foo$y.upd,foo$y.upd-1.96*sqrt(foo$P.upd),foo$y.upd+1.96*sqrt(foo$P.upd)),
        type="l",lty=1,col=c(1,2,2))
points(1:n,z,col=3)
lines(1:n,y,col=4)

#loglikelihood for different values of alpha
alpha.est = seq(0.8,0.999,length=100)
logl = rep(NA,length(alpha.est))
for(i in 1:length(alpha.est))
  logl[i] = kalman(c(logit(alpha.est[i]),log(sigma),log(tau)),z=z)$loglik
plot(alpha.est,logl,type="l")
theta0 = c(logit(alpha),log(sigma),log(tau))

#Optimize with respect to all parameters using the nlm routine
#Note that nlm MINIMIZE so that we need to change sign to get maximizationx
minusloglik = function(theta,z){-kalman(theta,z)$loglik}
res = nlm(minusloglik,theta0,z=z)
print(c(invlogit(res$est[1]),exp(res$est[2]),exp(res$est[3])))