library(MASS)
library(numDeriv)

## ----------------------------------------------------------------
# The falling body example from the book
## ----------------------------------------------------------------

## ----------------------------------------------------------------
# Parameters for simulation
A <- matrix(c(1,1,
              0,1), nrow=2, byrow=TRUE)
B <- matrix(c(-0.5,-1), nrow=2)
C <- matrix(c(1,0), nrow=1)
Sigma1 <- matrix(c(2,   0.8,
                   0.8, 1  ), nrow=2, byrow=TRUE)
Sigma2 <- 1000
# The state-vector init value
X0 <- c(10000, 0)
# Number of time points to simulate
n <- 100
# Vector for saving the output
y <- rep(NA, n)
# The input (i.e. gravity, hence constant input not changing over time)
g <- 9.82

# Run the simulation
X <- X0
for (i in 2:n){
    X <- A %*% X + B%*%g + chol(Sigma1) %*% matrix(rnorm(2),ncol=1)
    y[i] <- C %*% X + sqrt(Sigma2) %*% rnorm(1)
}

# Plot the simulated height
plot(y, type="l", xlab="time", ylab="Observed altitude [m]", ylim=range(y,na.rm=TRUE))

## ----------------------------------------------------------------
# The Kalman filter in a function returning the negative log-likelihood
negloglik <- function(prm){
  ## --------------------------------
  # Overwrite for the parameters given
  all[names(prm)] <- prm
  # Recreate the parameter values
  A <- matrix(all[c("A1","A2","A3","A4")], nrow=2)
  B <- matrix(all[c("B1","B2")], nrow=2)
  C <- matrix(all[c("C1","C2")], nrow=1)
  Sigma1 <- matrix(all[c("Sigma11","Sigma12","Sigma13","Sigma14")], nrow=2)
  Sigma2 <- matrix(all["Sigma2"])
  # Initial state vector and covariance
  X <- all[c("X01","X02")]
  SigmaX <- matrix(all[c("SigmaX01","SigmaX02","SigmaX03","SigmaX04")], nrow=2)
  ## --------------------------------
  # For keeping results
  n <- length(y)
  lik <- rep(NA,n)
  # Start on the iterations
  i <- 1
  while(i < n){
    # ---- PREDICT ----
    X <- A %*% X + B %*% g
    SigmaX <- A %*% SigmaX %*% t(A) + Sigma1
    SigmaY <- C %*% SigmaX %*% t(C) + Sigma2
    # ---- Step ----
    i <- i + 1
    # The likelihood
    lik[i] <- dnorm(y[i], mean=C%*%X, sd=sqrt(SigmaY))
    # ---- UPDATE ----
    K <- SigmaX %*% t(C) %*% solve(SigmaY)
    X <- X + K %*% (y[i] - C %*% X)
    SigmaX <- SigmaX - K %*% SigmaY %*% t(K)
  }
  # Calculate the negative log-likelihood
  return(-sum(log(lik[!is.na(lik)])))
}

# Prm in vector for estimation
prm <- c(
  B = c(-0.5,-1)
)
         
# All prm in vector (these will be the values if not over-writin from prm)
all <- c(
  A = c(A),
  B = c(B),
  C = c(C),
  X0 = X0,
  Sigma1 = c(Sigma1),
  Sigma2 = Sigma2,
  SigmaX0 = c(matrix(c(1000,   0,
                          0,1000), nrow=2, byrow=TRUE))
)

# Use an optimizer for estimation
(val <- optim(prm, negloglik, control=list(trace=0), method="L-BFGS-B"))

#----------------------------------------------------------------
# What about the estimated uncertainty of the parameters with MLE?
#----------------------------------------------------------------
# Approximate the Hessian
H <- hessian(negloglik, val$par)
H

# Fisher's information matrix
I <- qr.solve(H)

# Covariance matrix of the estimates
cov2cor(I)

# Standard error (Standard deviation estimate of the theta parameters, so mu and exp(sigma^2))
sqrt(diag(I))


#----------------------------------------------------------------
# Use the Kalman filter from the FKF package
library("FKF")
obj <- function(prm){
  # Overwrite for the parameters given
  all[names(prm)] <- prm
  # Recreate the parameter values
  A <- matrix(all[c("A1","A2","A3","A4")], nrow=2)
  B <- matrix(all[c("B1","B2")], nrow=2)
  C <- matrix(all[c("C1","C2")], nrow=1)
  Sigma1 <- matrix(all[c("Sigma11","Sigma12","Sigma13","Sigma14")], nrow=2)
  Sigma2 <- matrix(all["Sigma2"])
  # Initial state vector and covariance
  X0 <- all[c("X01","X02")]
  SigmaX0 <- matrix(all[c("SigmaX01","SigmaX02","SigmaX03","SigmaX04")], nrow=2)
  # Use the Kalman filter
  kf1 <- fkf(a0=X0, P0 = SigmaX0, dt = B%*%g, Tt = A, ct=0, Zt = C, HHt = Sigma1, GGt = Sigma2, yt = matrix(y,nrow=1))
  -kf1$logLik
}
# Estimate
(val <- optim(prm, obj, control=list(trace=0), method="L-BFGS-B"))


#----------------------------------------------------------------
# Now we should do a model validation!!
# Not implemented...

# E.g. plot one step predictions
plot(y, xlab="Frame", ylab="height [m]")
with(kf1, matlines((at[1,]) + cbind(0,-1.96*sqrt(Pt[1,1,]),1.96*sqrt(Pt[1,1,])),type="l", lty=c(1,2,2), col=2))

# White noise residuals ...