## Particle Filter EM algorithm.

## There are two steps to this algorithm.

## (Stochastic) E-Step We take <nparticle> draws from the competency
## trajectories, and weight them according to the likelihood under
## the current parameters, and calculate the weights for each trajectory

## M-Step, we maximize the Evidence Model parameters using weighted
## MLE.

## For the purposes of the M-step, the particles can be treated as
## additional replicates, and appended to the bottom of the matrix.

## The number of particles can be gradually increased as the number of
## cycles go by.  This may help with look maxima, and should cause
## earlier iterations to run faster.

## There are two variants on the algorithm depending on how the CGM
## parameters are estimated.  Type I uses EM to estimate CGM
## parameters as well as EM parameters.  Type II use PF techniques to
## estimate the CGM parameters, and Weighted ML to estimate the EM
## parameters.

###################################################
## Data Storage

## For the most part, we can treat cases and particles
## interchangably.  So if we have <ncases> cases and <nparticles>
## particles, we can store this in a matrix with <ncases>*<nparticle>
## rows.

## In a type I filter, the only place we need to worry about particles
## vs cases is when calculating the weights.  Even then, it might not
## be necessary, we can simply do weighted ML using each of the
## possible trajectories as weights.

## In the case of a type II filter, we do have to worry about
## particles as each particle corresponds to a different set of
## parameters.


PFEMData <-
setClass("PFEMData",
         representation(nparticles="integer",
                        obs="list",
                        latent="list",
                        background="list",
                        success = "list or NULL",
                        successObserved="logical",
                        weight="vector"))

setGeneric("ncases",function(pfdat) standardGeneric("ncases"))

## These functions return data replicated over particles.
setGeneric("getPFdatObs",function(pfdat,time) standardGeneric("getPFdatObs"))
setGeneric("getPFdatBack",function(pfdat,time) standardGeneric("getPFdatBack"))
setGeneric("getPFdatWeight",function(pfdat) standardGeneric("getPFdatWeight"))
setGeneric("getPFdatLatent",
           function(pfdat,time) standardGeneric("getPFdatLatent"))
setGeneric("getPFdatLatent<-",
           function(pfdat,time,value) standardGeneric("getPFdatLatent<-"))
setGeneric("getPFdatSuccess",
           function(pfdat,time) standardGeneric("getPFdatSuccess"))
setGeneric("getPFdatSuccess<-",
           function(pfdat,time,value) standardGeneric("getPFdatSuccess<-"))
## This function slices one particle of the latent variables.
setGeneric("getPFdatLatentPart",
           function(pfdat,iparticle,time)
           standardGeneric("getPFdatLatentPart"))
setGeneric("getPFdatLatentPart<-",
           function(pfdat,iparticle,time,value)
           standardGeneric("getPFdatLatentPart<-"))
setGeneric("getPFdatSuccessPart",
           function(pfdat,iparticle,time)
           standardGeneric("getPFdatSuccessPart"))
setGeneric("getPFdatSuccessPart<-",
           function(pfdat,iparticle,time,value)
           standardGeneric("getPFdatSuccessPart<-"))

setMethod("ncases","PFEMData",function(pfdat) nrow(pfdat@obs))

setMethod("getPFdatLatentPart","PFEMData",
          function (pfdat,iparticle,time) {
            ncases <- ncases(pfdat)
            if (is.logical(iparticle)) {
              selcases <- rep(iparticle, each=ncases)
            } else {
              selcases <-
                as.vector(outer(1:ncases,(iparticle-1)*ncases,"+"))
            }
            pfdat@latent[[time]][selcases,]
          })

setMethod("getPFdatLatentPart<-","PFEMData",
          function (pfdat,iparticle,time,value) {
            ncases <- ncases(pfdat)
            if (is.logical(iparticle)) {
              selcases <- rep(iparticle, each=ncases)
            } else {
              selcases <-
                as.vector(outer(1:ncases,(iparticle-1)*ncases,"+"))
            }
            pfdat@latent[[time]][selcases,] <- value
            pfdat
          })



setMethod("getPFdatLatent","PFEMData",
          function (pfdat,time) {
            pfdat@latent[[time]]
          })

setMethod("getPFdatLatent<-","PFEMData",
          function (pfdat,time,value) {
            pfdat@latent[time] <- value
            pfdat
          })

setMethod("getPFdatSuccessPart","PFEMData",
          function (pfdat,iparticle,time) {
            ncases <- ncases(pfdat)
            if (pfdat@successObserved) {
              selcases <- 1:ncases
            } else if (is.logical(iparticle)) {
              selcases <- rep(iparticle, each=ncases)
            } else {
              selcases <-
                as.vector(outer(1:ncases,(iparticle-1)*ncases,"+"))
            }
            pfdat@success[[time]][selcases,]
          })

setMethod("getPFdatSuccessPart<-","PFEMData",
          function (pfdat,iparticle,time,value) {
            ncases <- ncases(pfdat)
            if (pfdat@successObserved) {
              selcases <- 1:ncases
            } else if (is.logical(iparticle)) {
              selcases <- rep(iparticle, each=ncases)
            } else {
              selcases <-
                as.vector(outer(1:ncases,(iparticle-1)*ncases,"+"))
            }
            pfdat@success[[time]][selcases,] <- value
            pfdat
          })



setMethod("getPFdatSuccess","PFEMData",
          function (pfdat,time) {
            pfdat@success[[time]]
          })

setMethod("getPFdatSuccess<-","PFEMData",
          function (pfdat,time,value) {
            pfdat@success[[time]] <- value
            pfdat
          })

setMethod("getPFdatBack","PFEMData",
          function (pfdat,time) {
            if (length(pdfat@back) == 1) {
              back1 <- pfdat@back[[1]]
            } else {
              back1 <- pfdat@back[[time]]
            }
            d1 <- dim(back1)
            d1[1] <- d1[1]*pfdat@nparticles
            result <- array(NA,d1)
            cases <- 1:nrow(back1)
            for (iparticle in 1:pfdat@nparticles) {
              if (length(d1) ==2) {
                result[cases,] <- back1
              } else {
                result[cases,,] <- back1
              }
            }
            result
          })

setMethod("getPFdatObs","PFEMData",
          function (pfdat,time) {
            obs1 <- pfdat@obs[[time]]
            d1 <- dim(obs1)
            d1[1] <- d1[1]*pfdat@nparticles
            result <- array(NA,d1)
            cases <- 1:nrow(obs1)
            for (iparticle in 1:pfdat@nparticles) {
              if (length(d1) ==2) {
                result[cases,] <- obs1
              } else {
                result[cases,,] <- obs1
              }
            }
            result
          })

setMethod("getPFdatWeight","PFEMdata",
          function (pfdat) {
            rep(pfdat@weight,ncases(pfdat))
          })



############################################
## PF-Step

PFStepI <- function(pamodel,pfdata,nparticles) {
  ncycle <- pamodel
  ncases <- ncases(pfdata)
  pfdata@latent <- list()
  if (!pfdata@successObserved) {
    pdfata@success <- list()
  }

  ## Special Setup for cycle 1 (initial time point).
  latent <- drawInitialData(getCGM(pamodel),
                            getCGMparam(pamodel),
                            getPFdatBack(pfdata,icycle),
                            ncases*nparticles)
  getPFdatLatent(pfdata,1) <- latent
  weight <- llikeObs(getEMt(pamodel,1),
                     getEMparamt(pamodel,1),
                     getPFdatObs(pfdata,1),
                     latent,
                     getQt(pamodel,1),
                     getPFdatBack(pfdata,1))

  for (icycle in 2:ncycle) {

    ## Handle Success
    if (pfdata@successObserved) {
      success <- getPFdatSuccess(pfdata,icycle-1)
      weight <- weight +
        llikeSuccess(getCGM(pamodel),
                     getCGMparam(pamodel),
                     latent,
                     getPFdatAct(pdfata,icycle),
                     success,
                     getPFdatBack(pfdata,icycle))
    } else {
      success <- drawSuccess(getCGM(pamodel),
                             getCGMparam(pamodel),
                             latent,
                             getPFdatAct(pdfata,icycle),
                             getPFdatBack(pfdata,icycle))
      getPFdatSuccess(pfdata,icycle-1) <- success
    }

    latent <- advanceData(getCGM(pamodel),
                          getCGMparam(pamodel),
                          latent,
                          getPFdatAct(pdfata,icycle-1),
                          success,
                          getPFdatBack(pfdata,icycle-1),
                          getCycletime(pamodel,icycle-1))

    getPFdatLatent(pfdata,icycle) <- latent

    weight <- weight + llikeObs(getEMt(pamodel,icycle),
                                getEMparamt(pamodel,icycle),
                                getPFdatObs(pfdata,icycle),
                                latent,
                                getQt(pamodel,icycle),
                                getPFdatBack(pfdata,icycle))
  }
  pfdata@weight <- weight
  list(pamodel=pamodel,pfdata=pfdata)
}