### Various variants of the Kalman Filter.

## x[,,t] <- x[,,t-1]F[,,t] + u[,,t]B[,,t]+w[,,t]
## w[,,t] ~ N(0,Q[,,t])

## y[,,t] <- x[,,t]H[,,t] + v[,,t]
## v[,,t] ~ N(O,R[,,t])


## x[,,0] ~ N(mu,A)

## y == obs, x==est, u==act

#library(MCMCpack) ## For Wishart distribution
#library(mvtnorm) ## for Random normal

## Time homogenous Kalman Filter

TimeInvariantKalman <-
  setClass("TimeInvariantKalman",
  representation(fMean="matrix",fStd="matrix",
                 bMean="matrix",bStd="matrix",
                 hMean="matrix",hStd="matrix",
                 Qmean="matrix",Qdf="numeric",
                 Rmean="matrix",Rdf="numeric",
                 muMean="vector",muStd="matrix",
                 Amean="matrix",Adf="numeric"
                 ),
  prototype(fMean=matrix(c(.6,.3,.3,.6),2,2), fStd=matrix(.25,2,2),
            bMean=matrix(c(.2,.1,.1,.2),2,2), bStd=matrix(.05,2,2),
            hMean=matrix(c(.7,.3,.3,.7),2,2), hStd=matrix(.25,2,2),
            Qmean=matrix(c(.7,.3,.3,.7),2,2), Qdf=3,
            Rmean=matrix(c(.7,.3,.3,.7),2,2), Rdf=3,
            muMean=c(Mechanics=2,Fluency=2),  muStd=matrix(c(1,0,0,1),2,2),
            Amean=matrix(c(.6,.4,.4,.6),2,2), Adf=3,
            paramsFollowTime=FALSE,usesSuccess=FALSE,paramAdvances=FALSE),
  contains="BowtieFilter")

TimeInvariantKalmanParam <-
setClass("TimeInvariantKalmanParam",
         representation(Fmat="array",Bmat="array",Hmat="array",
                        Rinv="array",
                        Qmat="array",mu="matrix",Amat="array"))


## setMethod("initParam","TimeInvariantKalman",
## function(filter) {
##   Fmat <- array(rnorm(length(filter@fMean)*nparticles(filter)),
##                 c(dim(filter@fMean),nparticles(filter)))
##   Fmat <- sweep(sweep(Fmat,3,filter@fStd,"*"),3,filter@fMean,"+")


##   Bmat <- array(rnorm(length(filter@bMean)*nparticles(filter)),
##                 c(dim(filter@bMean),nparticles(filter)))
##   Bmat <- sweep(sweep(Bmat,3,filter@bStd,"*"),3,filter@bMean,"+")

##   Hmat <- array(rnorm(length(filter@hMean)*nparticles(filter)),
##                 c(dim(filter@hMean),nparticles(filter)))
##   Hmat <- sweep(sweep(Hmat,3,filter@hStd,"*"),3,filter@hMean,"+")

##   mu <- array(rnorm(length(filter@muMean)*nparticles(filter)),
##               c(length(filter@muMean),nparticles(filter)))
##   mu <- sweep(sweep(mu,2,filter@muStd,"*"),2,filter@muMean,"+")

##   Rinv <- array(0,c(dim(filter@Rmean),nparticles(filter)))
##   Qmat <- array(0,c(dim(filter@Qmean),nparticles(filter)))
##   Amat <- array(0,c(dim(filter@Amean),nparticles(filter)))
##   for (iparticle in 1:nparticles(filter)) {
##     Rinv[,,iparticle] <- rwish(filter@Rdf,filter@Rmean)
##     Qmat[,,iparticle] <- riwish(filter@Qdf,filter@Qmean)
##     Amat[,,iparticle] <- riwish(filter@Adf,filter@Amean)
##   }
##   new("TimeInvariantKalmanParam",
##       Fmat=Fmat, Bmat=Bmat, Hmat=Hmat, Rinv=Rinv,
##       Qmat=Qmat, mu=mu, Amat=Amat )
## })


## setMethod("nlatent","TimeInvariantKalman",
##           function(filter) {
##             length(filter@muMean)
##           })

## setMethod("nobserved","TimeInvariantKalman",
##           function(filter) {
##             ncol(filter@hMean)
##           })


## setMethod("naction","TimeInvariantKalman",
##           function(filter) {
##             nrow(filter@bMean)
##           })

## setMethod("nsuccess","TimeInvariantKalman",
##           function(filter) { 0 })



## TODO worry about names
## setMethod("initData",c("TimeInvariantKalman","TimeInvariantKalmanParam"),
## function(filter,param,ncases) {
##   result <- array(rnorm(ncases*nlatent(filter)*nparticles(filter)),
##                   c(ncases,nlatent(filter),nparticles(filter)))
##   for (iparticle in 1:nparticles(filter)) {
##     A <- chol(param@Amat[,,iparticle])
##     result[,,iparticle] <- sweep(result[,,iparticle]%*%A,2,
##                                  param@mu[,iparticle],"+")
##   }
##   result
## })

## setMethod("obsLlike",c("TimeInvariantKalman","TimeInvariantKalmanParam"),
## function (filter,param,obs,latent) {
##   ncases <- dim(obs)[1]
##   result <- rep(0,nparticles(filter))
##   for (iparticle in 1:nparticles(filter)) {
##     H <- param@Hmat[,,iparticle]
##     Rinv <- param@Rinv[,,iparticle]
##     v <- obs - latent[,,iparticle] %*% H
##     ker <- 0
##     for (i in 1:ncases) {
##       ker <- ker - v[i,]%*%Rinv%*%v[i,]
##     }
##     ## Do we need to rework this to use the log weights??  Yes!
##     result[iparticle] <- log(sqrt(det(Rinv))) + ker/2
##   }
##   result
## })

## setMethod("obsGen",c("TimeInvariantKalman","TimeInvariantKalmanParam"),
## function (filter,param,latent) {
##   ncases <- dim(latent)[1]
##   nobs <- nobserved(filter)
##   Zero <- rep(0,nobs)
##   result <- array(NA,c(ncases,nobs,nparticles(filter)))
##   for (iparticle in 1:nparticles(filter)) {
##     H <- param@Hmat[,,iparticle]
##     R <- solve(param@Rinv[,,iparticle])
##     v <- rmvnorm(ncases, Zero, R)
##     result[,,iparticle] <- latent[,,iparticle] %*% H + v
##   }
##   result
## })



## setMethod("advanceData",c("TimeInvariantKalman","TimeInvariantKalmanParam"),
## function (filter,param,latent,action,success,time) {
##   ncases <- dim(latent)[1]
##   nvar <- nlatent(filter)
##   result <- array(0,c(ncases,nvar,nparticles(filter)))
##   if (is.null(dim(action))) {
##     action <- matrix(action,ncases,nvar,byrow=TRUE)
##   }
##   for (iparticle in 1:nparticles(filter)) {
##     Fmat <- param@Fmat[,,iparticle]
##     Bmat <- param@Bmat[,,iparticle]
##     Q <- chol(param@Qmat[,,iparticle])
##     w <- matrix(rnorm(nvar*ncases),ncases,nvar)%*%Q
##     result[,,iparticle] <- latent[,,iparticle]%*%Fmat + action%*%Bmat + w
##   }
##   result
## })