library("coda")
## These are probably necessary, but whichever of the two is needed
## will almost certainly be loaded before the corresponding function
## is called.
##library("rstan")
##library("rjags")


### These are a set of routines for converting stan output to CODA output.
### There are a number of reasons to want to do this:
### (1) To be able to deal with output from Stan, JAGS and OpenBUGS in a uniform way.
### (2) Stan does not allow you to extend chains, this can be done
### easily in Stan (or once converted to the array format).
### (3) Stanfit object retain all of the data points, even auxiliary
### variables that are of no interest.
### (4) Stanfit objects cannot be dumped/restored.
### (5) as.array(Stanfit) generates an error if the chain did not
### converge which halts an automatic script, want to issue a warning
### and return NULL instead.
### (6) The MCArray object used in rjags is not as flexible as the
### mcmc.list object from coda


### Stan functions

## Bob Carpenter's one-liner (from Stan group):
## s <- mcmc.list(lapply(1:ncol(fit), function(x) mcmc(as.array(fit)[,x,])))

## fit-- stan fit object
## param -- name of parameter *required*
## K - number of components, set to 0 for level2 parameters
stan2coda <- function (fit,param,K=0) {
  tryCatch({
      mclist <- mcmc.list(lapply(1:ncol(fit),
                                 function (chain) {
                                   extract <- as.array(fit,pars=param)[,chain,]
                                   if (is.null(dim(extract))) {
                                     ## Scalar parameter
                                     extract <- matrix(extract,ncol=1,
                                                       dimnames=list(NULL,param))
                                   }
                                   mcmc(extract)}))
      ## The Stan as.array function takes care of names
      ## except for scalars, which are handled above, so we should be good.
      mclist},
      error=function (e) {
        ## Will get here if Stan run stopped.
        print(e)
        NULL
      })
}

## Sample idea, only now loop over MCMC objects, returning a vector of
## parameters.  This is closer to the rjags output, also slightly more
## convenient
## pars -- list of parameters to fetch
## K - number of components, set to 0 for level2 parameters
stan2coda2 <- function (fit,pars,K=0) {
  result <- lapply(pars,function(par) stan2coda(fit,par,K))
  names(result) <- pars
  result
}


### JAGS functions

## samples-- output of jags.samples
## param -- name of parameter *required*
## K - number of components, set to 0 for level2 parameters
jags2coda <- function (samples,param,K=0) {
  tryCatch({
    mcmc <- as.mcmc.list(samples[[param]])
      if (K==0) {                       #level 2
        if (nvar(mcmc)==1) { ## scalar
          varnames(mcmc) <- param
        } else { ## Vector (length K)
          varnames(mcmc) <- paste(param,"[",1:nvar(mcmc),"]",sep="")
        }
      } else {                          #level 1
        N <- nvar(mcmc)/K
        varnames(mcmc) <- paste(param,"[", rep(1:N,K),",",
                                           rep(1:K,each=N), "]",sep="")
      }
    mcmc},
    error=function (e) {
      ## Will get here if Stan run stopped.
      print(e)
      NULL
    })
}

## pars -- list of parameters to fetch
## K - number of components, set to 0 for level2 parameters
jags2coda2 <- function (samples,pars,K=0) {
  result <- lapply(pars,function(par) jags2coda(samples,par,K))
  names(result) <- pars
  result
}


## Does an rbind with two MCMC or MCMC.list objects.
mcmcjoin <- function (mcmc1,mcmc2) {
  if (is.mcmc.list(mcmc1)) {
    as.mcmc.list(lapply(1:length(mcmc1),
                        function(chain)
                          mcmcjoin(mcmc1[[chain]],mcmc2[[chain]])))
  } else {
    as.mcmc(rbind(mcmc1,mcmc2))
  }
}

## This combines all of the chains of an mcmclist into one big MC sample
combineChains <- function (mclist) {
  if (is.mcmc.list(mclist)) {
    mcmc <- mcmc(do.call("rbind",mclist))
    varnames(mcmc) <- varnames(mclist)
    mcmc
  } else {
    mclist
  }
}

## This function takes the matrix in a R by I*K format and reshapes it to a
## RxIxK array -- this is useful for Level 1 parameters
Karray <- function (mat,K) {
  array(mat,c(nrow(mat),ncol(mat)/K,K))
}

## This seems to run afoul of an R bug, if you try to assign
## dinames(x)[[2]] <- value, when value is a scalare and dimnames is
## NULL, you get an error.  It does not seem to be a problem when
## value has length greater than 1.  This screws up the jags2coda
## function when working with scalars.

"varnames<-" <- function (x, value) {
    if (is.mcmc(x)) {
      if (is.null(dimnames(x))) {
        dimnames(x) <- list(NULL,value)
      } else {
        dimnames(x)[[2]] <- value
      }
    }
    else if (is.mcmc.list(x))
        for (i in 1:nchain(x)) varnames(x[[i]]) <- value
    else stop("Not an mcmc or mcmc.list object")
    x
}