#######################################################################################################
#   KEM function
#   Minimal error checking is done.  You should run is.marssMLE(MLEobj) before calling this.
#   Maximization using an EM algorithm with Kalman filter
#######################################################################################################
MARSSkem <- function(MLEobj) {
  MODELobj <- MLEobj[["marss"]]
  # This is a core function and does not check if user specified a legal or solveable model.
  # y is MLEobj$marss$data with the missing values replaced by 0
  kf.x0 <- ifelse(MODELobj[["tinitx"]] == 1, "x10", "x00") # the initial conditions treatment "x00" x0 is at t=0 or "x01" x0 is at t=1
  # kf.x0=x00 prior is defined as being E[x(t=0)|y(t=0)]; xtt[0]=x0; Vtt[0]=V0
  # kf.x1=x10 prior is defined as being E[x(t=0)|y(t=0)]; xtt1[1]=x0; Vtt1[1]=V0

  # The model will be form = marss, so use base function for that form here
  constr.type <- describe_marss(MODELobj)
  # Check that model is allowed given the EM algorithm constaints; returns some info on the model structure
  if (MLEobj[["control"]][["trace"]] != -1) {
    errhead <- "\nErrors were caught in MARSSkemcheck \n"
    errmsg <- " Try using foo=MARSS(..., fit=FALSE), then summary(foo$model) to see what model you are trying to fit.\n"
    tmp <- MARSSkemcheck(MLEobj)
    if (!tmp$ok) {
      cat(c(errhead, tmp$msg, errmsg))
      stop("Stopped in MARSSkemcheck() due to specification problem(s).\n", call. = FALSE)
    }
  }

  # set up holders for warning messages
  msg <- NULL
  stop.msg <- NULL
  msg.kem <- NULL
  msg.kf <- NULL
  msg.conv <- NULL # error messages

  ## attach would be risky here since user might have one of these variables in their workspace
  y <- MODELobj[["data"]] # must have time going across columns
  d <- MODELobj[["free"]] # D or free matrix
  f <- MODELobj[["fixed"]] # f matrix
  inits <- MLEobj[["start"]]
  model.el <- attr(MODELobj, "par.names")
  model.dims <- attr(MODELobj, "model.dims")
  n <- model.dims[["data"]][1]
  TT <- model.dims[["data"]][2]
  m <- model.dims[["x"]][1]
  Id <- list(m = diag(1, m), n = diag(1, n))
  IIm <- diag(1, m) # identity matrices

  control <- MLEobj$control
  stopped.with.errors <- FALSE
  kf <- NULL
  condition.limit <- 1E10

  ## Set up MLE object for the iterations
  MLEobj.iter <- MLEobj
  MLEobj.iter$constr.type <- constr.type
  MLEobj.iter$par <- list()
  ## assign the starting parameter values; use fixed values where fixed otherwise use inits
  for (elem in model.el) MLEobj.iter$par[[elem]] <- inits[[elem]]

  ## make a list of time-varying and fixed parameters
  time.varying <- fixed <- list()
  for (elem in model.el) {
    if (is.fixed(d[[elem]])) {
      MLEobj.iter$par[[elem]] <- matrix(0, 0, 1)
      fixed[[elem]] <- TRUE
    } else {
      fixed[[elem]] <- FALSE
    }
    if (model.dims[[elem]][3] == 1) {
      time.varying[[elem]] <- FALSE
    } else {
      time.varying[[elem]] <- TRUE
    } # time-varying
  }
  # flags for whether a 0 was set on R or Q diagonals; determines whether various is.zero diagonal matrices are recomputed
  set.degen <- list(Q = FALSE, R = FALSE, V0 = FALSE)
  # define a couple constants that come up a lot
  # this is L%*V0%*%t(L)
  tmpV <- tcrossprod(parmat(MLEobj.iter, "L", t = 1)$L %*% parmat(MLEobj.iter, "V0", t = 1)$V0, parmat(MLEobj.iter, "L", t = 1)$L)
  IIzV0 <- diag(as.numeric(diag(tmpV) == 0), m)
  IImIIzV0 <- (IIm - IIzV0)

  ## zero out the missing values
  y[is.na(y)] <- 0

  ## Set up variable for debuging and diagnostics
  iter.record <- list(par = NULL, logLik = NULL)

  ################# The main EM loop which will run until tol reached or max.iter reached
  #######################################################################################

  # set up the convergence flags
  cvg <- 1 + control$abstol
  MLEobj.iter$logLik <- NA # start with no value
  # 72 means no info yet; 0 means converged
  MLEobj.iter$conv.test <- list(convergence = 72, messages = "No convergence testing performed.\n", not.converged.params = names(coef(MLEobj.iter, type = "vector")), converged.params = c())

  if (control$silent == 2) cat("EM iteration: ")
  for (iter in 1:(control$maxit + 1)) { #+1 so that the iter.record and kf are run for the last EM iteration
    if (control$silent == 2) cat(iter, " ")
    ################# E STEP Estimate states given U,Q,A,R,B,X0 via Kalman filter
    #####################################################################################
    kf.last <- kf
    kf <- MARSSkf(MLEobj.iter) # kf selects the function based on MLEobj$fun.kf
    if (!kf$ok) {
      if (control$trace > 0) {
        msg.kf <- c(msg.kf, paste("iter=", iter, " ", kf$errors))
      } else {
        msg.kf <- kf$errors
      }
      stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem() because numerical errors were generated in the Kalman filter.\n", sep = "")
      stopped.with.errors <- TRUE
      break
    }
    MLEobj.iter$kf <- kf
    MLEobj.iter$logLik <- kf$logLik
    if (control$demean.states) {
      xbar <- try(apply(cbind(kf$x0T, kf$xtT), 1, mean))
      MLEobj.iter$kf$xtT <- kf$xtT - xbar
      MLEobj.iter$kf$x0T <- kf$x0T - xbar
    }

    Ey <- MARSShatyt(MLEobj.iter)
    if (!Ey$ok) {
      if (control$trace > 0) {
        msg.kf <- c(msg.kf, paste("iter=", iter, " ", Ey$errors))
      } else {
        msg.kf <- Ey$errors
      }
      stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem() because numerical errors were generated in MARSShatyt\n", sep = "")
      stopped.with.errors <- TRUE
      break
    }
    MLEobj.iter$Ey <- Ey

    # This is a diagnostic line that checks if the solution is becoming unstable; likelike.old is set first at iter=1
    if (iter > 1 && is.finite(loglike.old) == TRUE && is.finite(MLEobj.iter$logLik) == TRUE) cvg <- MLEobj.iter$logLik - loglike.old
    if (iter > 2 & cvg < -sqrt(.Machine$double.eps)) {
      if (control$trace > 0) {
        msg.kem <- c(msg.kem, paste("iter=", iter, " LogLike DROPPED.  old=", loglike.old, " new=", MLEobj.iter$logLik, "\n", sep = ""))
      } else {
        msg.kem <- "MARSSkem: The soln became unstable and logLik DROPPED.\n"
      }
    }

    ################
    # Keep a record of the iterations for debugging and convergence diagnostics
    ################################################################
    if (control$trace > 0) { # if trace is on, keep the full record over all iterations
      iter.record$par <- rbind(iter.record$par, coef(MLEobj.iter, type = "vector"))
      iter.record$logLik <- c(iter.record$logLik, MLEobj.iter$logLik)
      if (!is.null(MLEobj.iter[["kf"]][["errors"]])) {
        msg.kf <- c(msg.kf, paste("iter=", iter, " ", kf$errors, sep = ""))
      }
      MLEobj.iter$iter.record <- iter.record
    } else { # Otherwise keep just last (control$conv.test.deltaT+1) iterations for diagnostics
      iter.record$par <- rbind(iter.record$par, coef(MLEobj.iter, type = "vector"))
      iter.record$logLik <- c(iter.record$logLik, MLEobj.iter$logLik)
      tmp.len <- dim(iter.record$par)[1]
      if (tmp.len > (control$conv.test.deltaT + 1)) {
        iter.record$par <- as.matrix(iter.record$par[(tmp.len - control$conv.test.deltaT):tmp.len, , drop = FALSE])
        iter.record$logLik <- iter.record$logLik[(tmp.len - control$conv.test.deltaT):tmp.len]
      }
      MLEobj.iter$iter.record <- iter.record
    }

    ################
    # Convergence Test
    ################################################################
    if (iter >= control$minit) { # then do convergence testing
      if (cvg >= 0 && cvg < control$abstol) {
        if (iter >= control$min.iter.conv.test) {
          MLEobj.iter$conv.test <- loglog.conv.test(iter.record, iter, deltaT = control$conv.test.deltaT, tol = control$conv.test.slope.tol)
          if (MLEobj.iter$conv.test$convergence != 1) break # 1=not converged, keep going; 0=converged; anything else=problem
        } else {
          MLEobj.iter$conv.test$convergence <- 3
        } # abstol reached log-log hasn't run yet because min.iter.cov.test not reached
      } else {
        MLEobj.iter$conv.test$convergence <- 4
      } # minit reached but not abstol
    }
    if (iter > control$maxit) {
      iter <- control$maxit
      break
    } # reset iter to maxit since needed to determine if stopped due to reaching maxit

    # Store loglike for comparison to new one after parameters are updated
    loglike.old <- MLEobj.iter$logLik

    # set the parameters at t=1
    par1 <- parmat(MLEobj.iter, t = 1)

    ################# M STEP update U,Q,A,R,B,X0 via ML given x(t) estimate
    # Update Q and R
    # Run Kalman smoother again to update the hidden states expectations
    # Update the other parameters

    ################################################################
    # Get new R subject to its constraints
    ################################################################
    # For the degen test, I require that d is a design matrix;
    if (control[["allow.degen"]]) {
      tmp <- degen.test("R", MLEobj.iter, iter) # this will test degeneracy and replace diags with 0s if needed
      MLEobj.iter <- tmp$MLEobj
      msg.kem <- c(msg.kem, tmp$msg)
      # update d and f because some of the R diagonals may have been set to 0
      if (tmp$set.degen) { # then some diagonals set to 0 so need to update values
        d$R <- MLEobj.iter$marss$free$R
        f$R <- MLEobj.iter$marss$fixed$R
        kf <- MLEobj.iter$kf
        Ey <- MLEobj.iter$Ey
        fixed$R <- is.fixed(d$R)
        set.degen$R <- TRUE # flag so that the identity matrices are redone
      }
    }

    # Now run the standard EM update equations
    if (!fixed[["R"]]) {
      sum1 <- t.dR.dR <- 0
      Z <- par1$Z
      A <- par1$A
      for (i in 1:TT) {
        if (time.varying[["Z"]] & i > 1) {
          Z <- parmat(MLEobj.iter, "Z", t = i)$Z
        }
        if (time.varying[["A"]] & i > 1) {
          A <- parmat(MLEobj.iter, "A", t = i)$A
        }
        if (i == 1 || time.varying[["R"]]) {
          dR <- sub3D(d[["R"]], t = i) # by def, i goes to TT if time-varying
          t.dR.dR <- t.dR.dR + crossprod(dR)
        }
        hatyt <- Ey[["ytT"]][, i, drop = FALSE]
        hatyxt <- sub3D(Ey[["yxtT"]], t = i)
        hatOt <- sub3D(Ey[["OtT"]], t = i)
        hatPt <- kf[["VtT"]][, , i] + tcrossprod(kf[["xtT"]][, i, drop = FALSE])
        hatxt <- kf[["xtT"]][, i, drop = FALSE]
        sum1a <- (hatOt - tcrossprod(hatyxt, Z) - tcrossprod(Z, hatyxt) - tcrossprod(hatyt, A) - tcrossprod(A, hatyt) +
          tcrossprod(Z %*% hatPt, Z) + tcrossprod(Z %*% hatxt, A) + tcrossprod(A, Z %*% hatxt) + tcrossprod(A)) # A%*%t.A
        sum1a <- symm(sum1a) # enforce symmetry function from MARSSkf
        sum1 <- sum1 + crossprod(dR, vec(sum1a))
      }
      if (time.varying[["R"]]) {
        if (length(t.dR.dR) == 1) {
          inv.dR <- 1 / t.dR.dR
        } else {
          inv.dR <- pcholinv(t.dR.dR)
        }
      } else {
        if (length(t.dR.dR) == 1) {
          inv.dR <- (1 / t.dR.dR) / TT
        } else {
          inv.dR <- pcholinv(t.dR.dR) / TT
        }
      }

      MLEobj.iter[["par"]][["R"]] <- inv.dR %*% sum1 # .03
      par1[["R"]] <- parmat(MLEobj.iter, "R", t = 1)$R

      # Start~~~~~~~~Error checking
      R <- par1$R # reset
      for (i in 1:model.dims[["R"]][3]) {
        if (time.varying$R & i > 1) R <- parmat(MLEobj.iter, "R", t = i)$R
        if (any(eigen(R, symmetric = TRUE, only.values = TRUE)$values < 0)) {
          stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem: solution became unstable. R update is not positive definite.\n", sep = "")
          stopped.with.errors <- TRUE
          break
        }
      }
      if (stopped.with.errors) break
      if (control$safe && !fixed[["R"]]) {
        new.kf <- rerun.kf("R", MLEobj.iter, iter)
        if (!new.kf$ok) {
          stopped.with.errors <- TRUE
          msg.kf <- c(msg.kf, new.kf$msg.kf)
          stop.msg <- new.kf$stop.msg
          break
        } else {
          kf <- new.kf$kf
          MLEobj.iter$kf <- kf
          MLEobj.iter$logLik <- kf$logLik
          Ey <- MARSShatyt(MLEobj.iter)
          MLEobj.iter$Ey <- Ey
          msg.kem <- c(msg.kem, new.kf$msg.kem)
        }
      }
    } # R not fixed

    ################
    # Get new Q subject to its constraints
    ################################################################
    # Start the testing for 0s along the diagonal of Q
    # For the degen test, I require that d is a design matrix;
    if (control[["allow.degen"]]) {
      tmp <- degen.test("Q", MLEobj.iter, iter) # this will test degeneracy and replace diags with 0s if needed
      MLEobj.iter <- tmp$MLEobj
      msg.kem <- c(msg.kem, tmp$msg)
      if (tmp$set.degen) {
        d$Q <- MLEobj.iter$marss$free$Q
        f$Q <- MLEobj.iter$marss$fixed$Q
        kf <- MLEobj.iter$kf
        Ey <- MLEobj.iter$Ey
        fixed$Q <- is.fixed(d$Q)
        set.degen$Q <- TRUE # flag so that the identity matrices are redone
      }
    }

    # Then do the regular EM update
    if (!fixed[["Q"]]) { # dim d$Q =0 or d$Q all zeros
      # If you treat x0 as at t=1 then
      # S00 = 0; S11 = 0; S10 = 0; X1 = 0; X0 = 0; TT.numer = TT-1
      # Otherwise if x0 is at t=0 follow Shumway and Stoffer (S&S2006 Eqn 6.67-69)

      dQ <- sub3D(d$Q, t = 1) # won't be all zeros due to !is.fixed
      B <- par1$B
      U <- par1$U
      if (kf.x0 == "x00") {
        TT.numer <- TT
        # IImIIzV0 = (IIm-IIz$V0[,,1]); IIzV0 = IIz$V0[,,1]
        X0 <- IImIIzV0 %*% kf$x0T + IIzV0 %*% par1$x0
        S00 <- kf$V0T + tcrossprod(X0)
        S10 <- kf$Vtt1T[, , 1] + tcrossprod(kf$xtT[, 1, drop = FALSE], X0) # where diag.V0=0 kf$Vtt1T[,,1]=0 since x at t-1 is fixed
        S11 <- kf$VtT[, , 1] + tcrossprod(kf$xtT[, 1, drop = FALSE])
        X1 <- kf$xtT[, 1, drop = FALSE]
        sum1a <- (S11 - tcrossprod(B, S10) - tcrossprod(S10, B) + tcrossprod(B %*% S00, B)
          - tcrossprod(U, X1) - tcrossprod(X1, U) + tcrossprod(U, B %*% X0) + tcrossprod(B %*% X0, U) + tcrossprod(U)) # U%*%t.U
        sum1a <- symm(sum1a) # symmetry function from MARSSkf
        sum1 <- crossprod(dQ, vec(sum1a))
        t.dQ.dQ <- crossprod(dQ)
      }
      if (kf.x0 == "x10") {
        sum1 <- 0
        t.dQ.dQ <- 0
        TT.numer <- TT - 1
        # Gharamani treatment of initial condition; the initial condition specifies x at t=1
      }

      for (i in 2:TT) {
        S00 <- kf[["VtT"]][, , i - 1] + tcrossprod(kf[["xtT"]][, i - 1, drop = FALSE]) # sum 2:T E(xt1T%*%t(xt1T))
        S10 <- kf[["Vtt1T"]][, , i] + tcrossprod(kf[["xtT"]][, i, drop = FALSE], kf[["xtT"]][, i - 1, drop = FALSE]) # sum 2:T E(xtT%*%t(xt1T))
        S11 <- kf[["VtT"]][, , i] + tcrossprod(kf[["xtT"]][, i, drop = FALSE]) # sum 2:T E(xtT%*%t(xt1T))
        X0 <- kf[["xtT"]][, i - 1, drop = FALSE] # sum 2:T E(xt1T)
        X1 <- kf[["xtT"]][, i, drop = FALSE] # sum 2:T E(xtT)
        if (time.varying[["B"]]) {
          B <- parmat(MLEobj.iter, "B", t = i)$B
        }
        if (time.varying[["U"]]) {
          U <- parmat(MLEobj.iter, "U", t = i)$U
        }
        if (time.varying[["Q"]]) {
          dQ <- sub3D(d[["Q"]], t = i)
          t.dQ.dQ <- t.dQ.dQ + crossprod(dQ)
        }
        sum1a <- (S11 - tcrossprod(B, S10) - tcrossprod(S10, B) + tcrossprod(B %*% S00, B)
          - tcrossprod(U, X1) - tcrossprod(X1, U) + tcrossprod(U, B %*% X0) + tcrossprod(B %*% X0, U) + tcrossprod(U))
        sum1a <- symm(sum1a) # enforce symmetry function from MARSSkf
        sum1 <- sum1 + crossprod(dQ, vec(sum1a))
      }

      # pcholinv because there might be all zero cols in dQ; won't equal matrix(0,1,1) since !is.fixed
      if (time.varying[["Q"]]) {
        if (length(t.dQ.dQ) == 1) {
          inv.dQ <- 1 / t.dQ.dQ
        } else {
          inv.dQ <- pcholinv(t.dQ.dQ)
        }
      } else {
        t.dQ.dQ <- crossprod(dQ) # t.dQ%*%dQ
        if (length(t.dQ.dQ) == 1) {
          inv.dQ <- (1 / t.dQ.dQ) / TT.numer
        } else {
          inv.dQ <- pcholinv(t.dQ.dQ) / TT.numer
        }
      }
      # 0 will appear in par where there are all 0 cols in d since inv.dQ will be 0 row/col there
      MLEobj.iter$par$Q <- inv.dQ %*% sum1
      par1$Q <- parmat(MLEobj.iter, "Q", t = 1)$Q

      # Start~~~~~~~~~~~~Error checking
      Q <- par1$Q # reset
      for (i in 1:max(dim(d$Q)[3], dim(f$Q)[3])) {
        if (time.varying$Q & i > 1) Q <- parmat(MLEobj.iter, "Q", t = i)$Q
        if (any(eigen(Q, symmetric = TRUE, only.values = TRUE)$values < 0)) {
          stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem: solution became unstable. Q update is not positive definite.\n", sep = "")
          stopped.with.errors <- TRUE
          break
        }
      }
      if (stopped.with.errors) break
      if (control$safe && !fixed[["Q"]]) {
        new.kf <- rerun.kf("Q", MLEobj.iter, iter)
        if (!new.kf$ok) {
          stopped.with.errors <- TRUE
          msg.kf <- c(msg.kf, new.kf$msg.kf)
          stop.msg <- new.kf$stop.msg
          break
        } else {
          kf <- new.kf$kf
          MLEobj.iter$kf <- kf
          MLEobj.iter$logLik <- kf$logLik
          Ey <- MARSShatyt(MLEobj.iter)
          MLEobj.iter$Ey <- Ey
          msg.kem <- c(msg.kem, new.kf$msg.kem)
        }
      }
    } # if Q not fixed

    # This code sets up the IIz and star (inverse) matrices needed
    # This only needs to be done at iter=1 or if Q, R or V0 might have changed
    if (iter == 1 || set.degen[["Q"]] || set.degen[["R"]] || !fixed[["Q"]] || !fixed[["R"]] || !fixed[["V0"]]) {
      # Set up the variance matrices needed for the degenerate case
      if (iter == 1) {
        IIz <- star <- list() # IIz location of 0s on diagonal
        elems <- c("Q", "R", "V0")
        for (elem in elems) { # set up the arrays; only needed at iter=1
          # var-cov = G Q t(G) and H R t(H)
          elem1 <- "L"
          if (elem == "Q") elem1 <- "G"
          if (elem == "R") elem1 <- "H"
          star[[elem]] <- IIz[[elem]] <- array(0, dim = c(model.dims[[elem1]][1], model.dims[[elem1]][1], max(model.dims[[elem]][3], model.dims[[elem1]][3])))
        }
      } else {
        elems <- c("Q", "R", "V0")[c((!fixed[["Q"]] | set.degen[["Q"]]), (!fixed[["R"]] | set.degen[["R"]]), !fixed[["V0"]])]
      }
      for (elem in elems) {
        elem1 <- "L"
        if (elem == "Q") elem1 <- "G"
        if (elem == "R") elem1 <- "H"
        thedim <- model.dims[[elem1]][1]
        maxT <- model.dims[[elem]][3] # Q, R, or V0
        maxT1 <- model.dims[[elem1]][3] # G, H, or L
        # star$elem is mathbb{elem}; the bolded Q, R, and L in section 4.4 and 5
        # move up so only done once; star[[elem]]=IIz[[elem]]=array(0,dim=c(thedim,thedim,maxT))
        if (maxT == 1) QRV <- parmat(MLEobj.iter, elem, t = 1)[[elem]]
        if (maxT1 == 1) GHL <- parmat(MLEobj.iter, elem1, t = 1)[[elem1]]
        for (i in 1:max(maxT, maxT1)) {
          if (maxT != 1) QRV <- parmat(MLEobj.iter, elem, t = i)[[elem]]
          if (maxT1 != 1) GHL <- parmat(MLEobj.iter, elem1, t = i)[[elem1]]
          # this is being done to find the zeros on the diagonal
          pari <- tcrossprod(GHL %*% QRV, GHL)
          # These are the identity matrices used to identify the location of deterministic rows of x;
          # section 7.2 in EM Derivations "Idntifying the fully deterministic x rows"
          # IIz means location of 0s on diagonal of var-cov matrix GHL%*%QRV%*%t(GHL)
          if (iter == 1 || set.degen[[elem]]) {
            IIz[[elem]][, , i] <- diag(as.numeric(diag(pari) == 0), thedim)
            # I the locations of 0s on diagonal of Q are time-constant; see section 7.2
            if (elem == "Q" & max(maxT, maxT1) != 1) {
              if (!all.equal(IIz[[elem]][, , 1], IIz[[elem]][, , i])) {
                stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem. IIz$Q (location of 0s on diagonal) must be time invariant.\nYou probably want to set allow.degen=FALSE if it is true.\n", sep = "")
                stopped.with.errors <- TRUE
                break
              }
            }
          }
          # this is defining the bolded Q, R and K in section 4.4 and 5 of the EM Derivation
          starmultiplier <- tcrossprod(pcholinv(crossprod(GHL)), GHL)
          star[[elem]][, , i] <- crossprod(starmultiplier, pcholinv(QRV) %*% starmultiplier)
        }
        set.degen[[elem]] <- FALSE # reset so this code is not run again
        if (elem == "V0" && (iter == 1 || set.degen[["V0"]])) { # then 0 location in V0 has potentially changed (via degen.test(V0))
          # set up the diag matrices needed often
          IIzV0 <- sub3D(IIz[["V0"]], t = 1)
          IImIIzV0 <- (IIm - IIzV0)
        }
      } # for over elems
      if (stopped.with.errors) break
    } # if Q, R or V0 is not fixed or a value was set to 0 via allow.degen or iter=1

    ################################################################
    # Get new x0 subject to its constraints
    ################################################################
    if (!fixed[["x0"]]) { # some element needs estimating
      f.x0 <- sub3D(f$x0, t = 1)
      d.x0 <- sub3D(d$x0, t = 1)
      A <- par1$A
      Z <- par1$Z
      B <- par1$B
      U <- par1$U
      Qinv <- sub3D(star$Q, t = 1)
      diag.Q <- 1 - takediag(IIz$Q[, , 1])
      Rinv <- sub3D(star$R, t = 1)
      diag.R1 <- 1 - takediag(IIz$R[, , 1])
      IIz.R <- sub3D(IIz$R, t = 1)
      x0.degen.update <- FALSE
      diag.V0 <- 1 - takediag(IIzV0)
      nQ0 <- sum(diag.Q == 0)
      if (any(diag.Q == 0)) x0.degen.update <- !is.fixed(d.x0[diag.Q == 0, , drop = FALSE])
      numer <- denom <- 0
      if (kf.x0 == "x00") hatxt0 <- kf$x0T else hatxt0 <- kf$xtT[, 1, drop = FALSE]
      if (any(diag.V0 == 1)) { # meaning some V0 positive
        denom <- crossprod(d.x0, star$V0[, , 1] %*% d.x0)
        numer <- crossprod(d.x0, hatxt0 - f.x0)
      }

      if (any(diag.V0 == 0)) {
        AdjM <- B
        AdjM[AdjM != 0] <- 1
        if (kf.x0 == "x00") {
          # set up values for t=1
          Bstar <- B
          Bstar.tm <- IIm
          Ustar <- U
          Ustar.tm <- 0 * U
          Mt <- AdjM
          IId.tm <- IIm # t=0; no w's
          IId <- makediag(1 - diag.Q) # only can be w free if Q==0
          if (any(diag.Q == 0) & any(diag.Q != 0)) { # which did not pick up a zero from B
            IId[diag.Q == 0, diag.Q == 0][1 + 0:(nQ0 - 1) * (nQ0 + 1)] <- apply(Mt[diag.Q == 0, diag.Q != 0, drop = FALSE] == 0, 1, all)
          }
          Delta7 <- kf$xtT[, 1, drop = FALSE] - B %*% (IImIIzV0 %*% hatxt0 + IIzV0 %*% f.x0) - U
          Delta8 <- B %*% IIzV0 %*% d.x0 # since IId.tm and Bstar.tm are IIm
          numer <- numer + crossprod(Delta8, Qinv %*% Delta7)
          denom <- denom + crossprod(Delta8, Qinv %*% Delta8)
        } else { # x10
          Bstar <- IIm
          Ustar <- 0 * U
          IId.tm <- 0
          IId <- IIm
          Mt <- IIm
        }
        if (any(IId == 1)) { # again t=1
          Delta5 <- Ey$ytT[, 1, drop = FALSE] - Z %*% ((IIm - IId) %*% kf$xtT[, 1, drop = FALSE]) - Z %*% IId %*% (Bstar %*% (IImIIzV0 %*% hatxt0 + IIzV0 %*% f.x0) + Ustar) - A
          Delta6 <- Z %*% IId %*% Bstar %*% (IIzV0 %*% d.x0)
          if (any(diag.R1 == 0)) {
            if (any(crossprod(Delta6, IIz.R) %*% Delta6 != 0)) {
              stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem at x0 update.\n There are 0s on R diagonal. x0 assoc with these must be fixed (not estimated)\n when using the EM algorithm. Try method=\"BFGS\".  Type MARSSinfo(\"x0R0\") for help.\n", sep = "")
              stopped.with.errors <- TRUE
              break
            }
          }
          numer <- numer + crossprod(Delta6, Rinv %*% Delta5)
          denom <- denom + crossprod(Delta6, Rinv %*% Delta6)
        }
        # t>1 will break out as soon as no IId=1
        for (t in 2:TT) { # start at 2 if x00; at 3 if x10
          if (!any(IId == 1)) break
          if (time.varying$A) A <- parmat(MLEobj.iter, c("A"), t = t)$A
          if (time.varying$B) B <- parmat(MLEobj.iter, c("B"), t = t)$B
          if (time.varying$U) U <- parmat(MLEobj.iter, c("U"), t = t)$U
          if (time.varying$Z) Z <- parmat(MLEobj.iter, c("Z"), t = t)$Z
          if (time.varying$R) {
            Rinv <- star$R[, , t]
            IIz.R <- sub3D(IIz$R, t = t)
          }
          if (time.varying$Q) Qinv <- star$Q[, , t]
          Ustar.tm <- Ustar
          Ustar <- B %*% Ustar + U
          Bstar.tm <- Bstar
          Bstar <- B %*% Bstar
          if (t <= (m + 1)) {
            IId.tm <- IId
            Mt <- AdjM %*% Mt
            IId <- makediag(1 - diag.Q) # only can be w free if Q==0
            if (any(diag.Q == 0) & any(diag.Q != 0)) {
              IId[diag.Q == 0, diag.Q == 0][1 + 0:(nQ0 - 1) * (nQ0 + 1)] <- apply(Mt[diag.Q == 0, diag.Q != 0, drop = FALSE] == 0, 1, all)
            }
          }
          if (any(IId == 1)) {
            Delta5 <- Ey$ytT[, t, drop = FALSE] - Z %*% ((IIm - IId) %*% kf$xtT[, t, drop = FALSE]) - Z %*% IId %*% (Bstar %*% (IImIIzV0 %*% hatxt0 + IIzV0 %*% f.x0) + Ustar) - A
            Delta6 <- Z %*% IId %*% Bstar %*% (IIzV0 %*% d.x0)
            # Deal with Delta6=0 and Rinv=Inf, so 0*Inf
            if (any(diag.R1 == 0)) {
              if (any(crossprod(Delta6, IIz.R) %*% Delta6 != 0)) {
                stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem at x0 update.\n There are 0s on R diagonal. x0 assoc with these must be fixed (not estimated)\n when using the EM algorithm. Try method=\"BFGS\".  Type MARSSinfo(\"x0R0\") for help.\n", sep = "")
                stopped.with.errors <- TRUE
                break
              }
            }
            numer <- numer + crossprod(Delta6, Rinv %*% Delta5)
            denom <- denom + crossprod(Delta6, Rinv %*% Delta6)
          }
          if (any(IId.tm == 1)) {
            Delta7 <- kf$xtT[, t, drop = FALSE] - B %*% (IIm - IId.tm) %*% kf$xtT[, t - 1, drop = FALSE] - B %*% IId.tm %*% (Bstar.tm %*% (IImIIzV0 %*% hatxt0 + IIzV0 %*% f.x0) + Ustar.tm) - U
            Delta8 <- B %*% IId.tm %*% Bstar.tm %*% (IIzV0 %*% d.x0)
            numer <- numer + crossprod(Delta8, Qinv %*% Delta7)
            denom <- denom + crossprod(Delta8, Qinv %*% Delta8)
          }
        } # for t
      } # any diag.LAM=0
      if (length(denom) == 1) {
        denom <- 1 / denom
      } else {
        denom <- try(pcholinv(denom))
      }
      if (inherits(denom, "try-error") || (length(denom) == 1 && denom == Inf)) {
        stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem at x0 update. denom is not invertible. \n This means that some of the x0 cannot be estimated. Type MARSSinfo('denominv') for more info. \n", sep = "")
        stopped.with.errors <- TRUE
        break
      }
      MLEobj.iter$par$x0 <- denom %*% numer
      if (!is.matrix(MLEobj.iter$par$x0)) MLEobj.iter$par$x0 <- matrix(MLEobj.iter$par$x0, dim(d$x0)[2], 1)
      par1$x0 <- parmat(MLEobj.iter, "x0", t = 1)$x0

      # ~~~~~~~~Error checking
      if (control$safe && !fixed[["x0"]]) {
        new.kf <- rerun.kf("x0", MLEobj.iter, iter)
        if (!new.kf$ok) {
          stopped.with.errors <- TRUE
          msg.kf <- c(msg.kf, new.kf$msg.kf)
          stop.msg <- new.kf$stop.msg
          break
        } else {
          kf <- new.kf$kf
          MLEobj.iter$kf <- kf
          MLEobj.iter$logLik <- kf$logLik
          Ey <- MARSShatyt(MLEobj.iter)
          MLEobj.iter$Ey <- Ey
          msg.kem <- c(msg.kem, new.kf$msg.kem)
        }
      }
    } # x0 is not fixed

    ################
    # Get new V0 subject to its constraints
    ################################################################
    if (!fixed[["V0"]]) { # some element needs estimating (obviously V0!=0)
      dV0 <- sub3D(d$V0, t = 1)
      V0.update <- chol2inv(chol(crossprod(dV0))) %*% crossprod(dV0, vec(kf$V0T))
      MLEobj.iter$par$V0 <- V0.update
      if (!is.matrix(MLEobj.iter$par$V0)) MLEobj.iter$par$V0 <- matrix(MLEobj.iter$par$V0, dim(d$V0)[2], 1)
      par1$V0 <- parmat(MLEobj.iter, "V0", t = 1)$V0

      # ~~~~~~~~Error checking
      if (any(eigen(par1$V0, symmetric = TRUE, only.values = TRUE)$values < 0)) {
        tmp <- ""
        if (any(abs(Re(eigen(par1$B, only.values = TRUE)$values) > 1))) tmp <- "Your B matrix is outside the unit circle.  This is likely the problem.\n"
        stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem: solution became unstable. V0 update is not positive definite.\n", tmp, sep = "")
        stopped.with.errors <- TRUE
        break
      }

      if (control$safe) {
        new.kf <- rerun.kf("V0", MLEobj.iter, iter)
        if (!new.kf$ok) {
          stopped.with.errors <- TRUE
          msg.kf <- c(msg.kf, new.kf$msg.kf)
          stop.msg <- new.kf$stop.msg
          break
        } else {
          kf <- new.kf$kf
          MLEobj.iter$kf <- kf
          MLEobj.iter$logLik <- kf$logLik
          Ey <- MARSShatyt(MLEobj.iter)
          MLEobj.iter$Ey <- Ey
          msg.kem <- c(msg.kem, new.kf$msg.kem)
        }
      }
    } # if not fixed V0

    ################
    # Get new A subject to its constraints (update of R will use this)
    ##############################################################
    if (!fixed[["A"]]) { # if there is anything to update
      # note if Z and f.a are constant then we can write this as numer = (Ey$ytT - Z %*% kf$xtT)%*%matrix(1,dim(kf$xtT)[2],1)-TT*f$A
      numer <- denom <- 0
      Z <- par1$Z # reset
      starR <- star[["R"]][, , 1]
      for (i in 1:TT) {
        if (time.varying[["Z"]] & i > 1) Z <- parmat(MLEobj.iter, "Z", t = i)$Z
        if (i == 1 || time.varying[["A"]]) {
          dA <- sub3D(d[["A"]], t = i)
          fA <- sub3D(f[["A"]], t = i)
        }
        if (time.varying[["R"]]) starR <- star[["R"]][, , i]
        numer <- numer + crossprod(dA, starR %*% (Ey[["ytT"]][, i, drop = FALSE] - Z %*% kf[["xtT"]][, i, drop = FALSE] - fA))
        denom <- denom + crossprod(dA, starR %*% dA)
      }
      if (length(denom) == 1) {
        denom <- try(1 / denom)
      } else {
        denom <- try(chol2inv(chol(denom)))
      }
      if (inherits(denom, "try-error")) {
        stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem at A update. denom is not invertible. \n If some of your R diagonals equal 0, then A elements corresponding to R==0 cannot be estimated.\n The problem may be with your D matrix (if you have one) also. Type MARSSinfo('denominv') for more info. \n", sep = "")
        stopped.with.errors <- TRUE
        break
      }
      MLEobj.iter$par$A <- denom %*% numer

      # not sure why denom%*%numer would ever not be a matrix
      if (!is.matrix(MLEobj.iter$par$A)) MLEobj.iter$par$A <- matrix(MLEobj.iter$par$A, dim(d$A)[2], 1)
      par1$A <- parmat(MLEobj.iter, "A", t = 1)$A

      # ~~~~~~~~Error checking
      if (control$safe) {
        new.kf <- rerun.kf("A", MLEobj.iter, iter)
        if (!new.kf$ok) {
          stopped.with.errors <- TRUE
          msg.kf <- c(msg.kf, new.kf$msg.kf)
          stop.msg <- new.kf$stop.msg
          break
        } else {
          kf <- new.kf$kf
          MLEobj.iter$kf <- kf
          MLEobj.iter$logLik <- kf$logLik
          Ey <- MARSShatyt(MLEobj.iter)
          MLEobj.iter$Ey <- Ey
          msg.kem <- c(msg.kem, new.kf$msg.kem)
        }
      }
    } # A not fixed

    ################
    # Get new U subject to its constraints (update of Q and B will use this)
    ########################################################################
    if (!fixed[["U"]]) { # if there is anything to update
      numer <- matrix(0, m, 1)
      denom <- matrix(0, m, m) # this is the start if kf.x0="x10"
      fU <- sub3D(f$U, t = 1)
      dU <- sub3D(d$U, t = 1)
      B <- par1$B
      Z <- par1$Z
      A <- par1$A # reset
      Qinv <- star$Q[, , 1]
      Rinv <- star$R[, , 1]
      # U.degen.update = FALSE #CUT?
      diag.Q <- 1 - takediag(IIz$Q[, , 1])
      diag.V0 <- 1 - takediag(IIzV0)
      nQ0 <- sum(diag.Q == 0)
      # if(any(diag.Q==0)) U.degen.update=!all(d$U[diag.Q==0,,]==0)    #CUT?
      numer <- denom <- 0
      if (kf.x0 == "x00") hatxt0 <- kf$x0T else hatxt0 <- kf$xtT[, 1, drop = FALSE]
      E.x0 <- IImIIzV0 %*% hatxt0 + IIzV0 %*% par1$x0
      AdjM <- B
      AdjM[AdjM != 0] <- 1
      if (kf.x0 == "x00") {
        # set up values for t=1
        Bstar <- B
        Bstar.tm <- IIm
        fstar <- fU
        fstar.tm <- 0 * fU
        Dstar <- dU
        Dstar.tm <- 0 * dU
        Mt <- AdjM
        IId.tm <- IIm # t=0; no w's
        IId <- makediag(1 - diag.Q) # only can be w free if Q==0
        if (any(diag.Q == 0) & any(diag.Q != 0)) { # which did not pick up a zero from B
          IId[diag.Q == 0, diag.Q == 0][1 + 0:(nQ0 - 1) * (nQ0 + 1)] <- apply(Mt[diag.Q == 0, diag.Q != 0, drop = FALSE] == 0, 1, all)
        }
        # x_1-B(I-Id)xtm-B Id (B* E.x0 + f*)-fU; f*=0, B*=I; xtm=E.x0 so reduces to the following
        Delta3 <- kf$xtT[, 1, drop = FALSE] - B %*% E.x0 - fU
        Delta4 <- dU # since IId.tm and Bstar.tm are IIm
        numer <- numer + crossprod(Delta4, Qinv %*% Delta3)
        denom <- denom + crossprod(Delta4, Qinv %*% Delta4)
      } else { # x10
        Bstar <- IIm
        fstar <- 0 * fU
        Dstar <- 0 * dU
        IId.tm <- 0 * IIm
        IId <- IIm
        Mt <- IIm
      }
      if (any(IId == 1)) { # again t=1
        Delta1 <- Ey$ytT[, 1, drop = FALSE] - Z %*% ((IIm - IId) %*% kf$xtT[, 1, drop = FALSE]) - Z %*% (IId %*% (Bstar %*% E.x0 + fstar)) - A
        Delta2 <- Z %*% IId %*% Dstar
        numer <- numer + crossprod(Delta2, Rinv %*% Delta1)
        denom <- denom + crossprod(Delta2, Rinv %*% Delta2)
      }
      for (t in 2:TT) {
        if (time.varying$U) {
          fU <- sub3D(f$U, t = t)
          dU <- sub3D(d$U, t = t)
        }
        if (time.varying$B) B <- parmat(MLEobj.iter, c("B"), t = t)$B
        if (time.varying$A) A <- parmat(MLEobj.iter, c("A"), t = t)$A
        if (time.varying$Z) Z <- parmat(MLEobj.iter, c("Z"), t = t)$Z
        if (time.varying$R) Rinv <- star$R[, , t]
        if (time.varying$Q) Qinv <- star$Q[, , t]
        fstar.tm <- fstar
        fstar <- B %*% fstar + fU
        Dstar.tm <- Dstar
        Dstar <- B %*% Dstar + dU
        Bstar.tm <- Bstar
        Bstar <- B %*% Bstar
        if (t <= (m + 1)) {
          IId.tm <- IId
          IId <- makediag(1 - diag.Q) # only can be w free if Q==0
          if (any(diag.Q == 0) & any(diag.Q != 0)) {
            Mt <- AdjM %*% Mt
            IId[diag.Q == 0, diag.Q == 0][1 + 0:(nQ0 - 1) * (nQ0 + 1)] <- apply(Mt[diag.Q == 0, diag.Q != 0, drop = FALSE] == 0, 1, all)
          }
        }
        if (any(IId == 1)) {
          Delta1 <- Ey$ytT[, t, drop = FALSE] - Z %*% ((IIm - IId) %*% kf$xtT[, t, drop = FALSE]) - Z %*% (IId %*% (Bstar %*% E.x0 + fstar)) - A
          Delta2 <- Z %*% IId %*% Dstar
          numer <- numer + crossprod(Delta2, Rinv %*% Delta1)
          denom <- denom + crossprod(Delta2, Rinv %*% Delta2)
        }
        Delta3 <- kf$xtT[, t, drop = FALSE] - B %*% ((IIm - IId.tm) %*% kf$xtT[, t - 1, drop = FALSE]) - B %*% (IId.tm %*% (Bstar.tm %*% E.x0 + fstar.tm)) - fU
        Delta4 <- dU + B %*% IId.tm %*% Dstar.tm
        numer <- numer + crossprod(Delta4, Qinv %*% Delta3)
        denom <- denom + crossprod(Delta4, Qinv %*% Delta4)
      } # for i
      if (length(denom) == 1) {
        denom <- 1 / denom
      } else {
        denom <- try(chol2inv(chol(denom)), silent = TRUE)
      }
      if (inherits(denom, "try-error") || (length(denom) == 1 && denom == 0)) {
        stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem at U update. denom is not invertible.\n This means some of the U (+ C) terms cannot be estimated.\n Type MARSSinfo('denominv') for more info. \n", sep = "")
        stopped.with.errors <- TRUE
        break
      }
      MLEobj.iter$par$U <- denom %*% numer
      if (!is.matrix(MLEobj.iter$par$U)) MLEobj.iter$par$U <- matrix(MLEobj.iter$par$U, dim(d$U)[2], 1)
      par1$U <- parmat(MLEobj.iter, "U", t = 1)$U

      # ~~~~~~~~Error checking
      if (control$safe && !fixed[["U"]]) {
        new.kf <- rerun.kf("U", MLEobj.iter, iter)
        if (!new.kf$ok) {
          stopped.with.errors <- TRUE
          msg.kf <- c(msg.kf, new.kf$msg.kf)
          stop.msg <- new.kf$stop.msg
          break
        } else {
          kf <- new.kf$kf
          MLEobj.iter$kf <- kf
          MLEobj.iter$logLik <- kf$logLik
          Ey <- MARSShatyt(MLEobj.iter)
          MLEobj.iter$Ey <- Ey
          msg.kem <- c(msg.kem, new.kf$msg.kem)
        }
      }
    } # any U not fixed

    ################
    # Get new B subject to its constraints
    ################################################################
    if (!fixed[["B"]]) {
      # t.kf.xtT = t(kf$xtT) #move t() out of for loop
      # need these for t=1 whether kf.x0=x00 or not
      U <- par1$U # reset
      starQ <- sub3D(star$Q, t = 1)
      dB <- sub3D(d$B, t = 1)
      fB <- sub3D(f$B, t = 1)
      if (kf.x0 == "x00") { # prior is defined as being E[x(t=0)|y(t=0)]; xtt[0]=x0; Vtt[0]=V0
        hatxtm <- IImIIzV0 %*% kf$x0T + IIzV0 %*% par1$x0
        hatVtm <- IImIIzV0 %*% kf$V0T %*% IImIIzV0 + IIzV0 %*% par1$V0 %*% IIzV0
        hatxt <- kf$xtT[, 1, drop = FALSE]
        Ptm <- hatVtm + tcrossprod(hatxtm, kf$x0T) # note def of t.hatxtm = kf$x0T not t(hatxtm)
        Pttm <- kf$Vtt1T[, , 1] + tcrossprod(hatxt, kf$x0T) # note def of t.hatstm for t=0
        kronPtmQ <- kronecker(Ptm, starQ)
        denom <- crossprod(dB, kronPtmQ %*% dB)
        numer <- crossprod(dB, (vec(starQ %*% (Pttm - tcrossprod(U, kf$x0T))) - kronPtmQ %*% fB))
      } else { # prior is defined as being E[x(t=1)|y(t=0)]; xtt1[1]=x0; Vtt1[1]=V0
        denom <- numer <- 0 # see Ghahramani and Hinton treatment.  Summation starts at t=2
      }

      for (i in 2:TT) {
        hatxtm <- kf$xtT[, i - 1, drop = FALSE]
        hatVtm <- kf$VtT[, , i - 1]
        hatxt <- kf$xtT[, i, drop = FALSE]
        Ptm <- hatVtm + tcrossprod(hatxtm)
        Pttm <- kf$Vtt1T[, , i] + tcrossprod(hatxt, hatxtm)
        if (time.varying$Q) starQ <- sub3D(star$Q, t = i)
        kronPtmQ <- kronecker(Ptm, starQ)
        if (time.varying$B) {
          dB <- sub3D(d$B, t = i)
          fB <- sub3D(f$B, t = i)
        }
        if (time.varying$U) U <- parmat(MLEobj.iter, "U", t = i)$U
        denom <- denom + crossprod(dB, kronPtmQ %*% dB)
        numer <- numer + crossprod(dB, vec(starQ %*% (Pttm - tcrossprod(U, hatxtm))) - kronPtmQ %*% fB)
      } # for i
      if (length(denom) == 1) {
        denom <- try(1 / denom)
      } else {
        denom <- try(chol2inv(chol(denom)))
      }
      if (inherits(denom, "try-error")) {
        stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem at B update. denom is not invertible.\n Type MARSSinfo('denominv') for more info. \n", sep = "")
        stopped.with.errors <- TRUE
        break
      }
      MLEobj.iter$par$B <- denom %*% numer
      if (!is.matrix(MLEobj.iter$par$B)) MLEobj.iter$par$B <- matrix(MLEobj.iter$par$B, dim(d$B)[2], 1)
      par1$B <- parmat(MLEobj.iter, "B", t = 1)$B

      # ~~~~~~~~Error checking
      if (control$safe) {
        new.kf <- rerun.kf("B", MLEobj.iter, iter)
        if (!new.kf$ok) {
          stopped.with.errors <- TRUE
          msg.kf <- c(msg.kf, new.kf$msg.kf)
          stop.msg <- new.kf$stop.msg
          break
        } else {
          kf <- new.kf$kf
          MLEobj.iter$kf <- kf
          MLEobj.iter$logLik <- kf$logLik
          Ey <- MARSShatyt(MLEobj.iter)
          MLEobj.iter$Ey <- Ey
          msg.kem <- c(msg.kem, new.kf$msg.kem)
        }
      }
      if (control$trace > 0) {
        Ck <- kappa(denom)
        if (Ck > condition.limit) msg.kem <- c(msg.kem, paste("iter=", iter, " Unstable B estimate because P_{t-1,t-1} is ill-conditioned. C =", round(Ck), "\n", sep = ""))
        for (i in 1:max(dim(f$B)[3], dim(d$B)[3])) {
          parB <- parmat(MLEobj.iter, "B", t = i)$B
          if (any(abs(Re(eigen(parB, only.values = TRUE)$values)) > 1)) msg.kem <- c(msg.kem, paste("iter=", iter, ",t=", i, " B update is outside the unit circle.", "\n", sep = ""))
        }
      }
    } # if !is.fixed B

    ################
    # Get new Z subject to its constraints
    ################################################################
    if (!fixed[["Z"]]) {
      numer <- denom <- 0
      A <- par1$A # reset
      starR <- star$R[, , 1]
      for (i in 1:TT) {
        hatxt <- kf$xtT[, i, drop = FALSE]
        Pt <- kf$VtT[, , i] + tcrossprod(hatxt)
        hatyxt <- Ey$yxtT[, , i]
        if (time.varying$A & i > 1) A <- parmat(MLEobj.iter, "A", t = i)$A
        if (time.varying$R) starR <- star$R[, , i]
        if (i == 1 || time.varying$Z) {
          fZ <- sub3D(f$Z, t = i)
          dZ <- sub3D(d$Z, t = i)
        }
        PkronR <- kronecker(Pt, starR)
        numer <- numer + crossprod(dZ, vec(starR %*% (hatyxt - tcrossprod(A, hatxt))) - PkronR %*% fZ)
        denom <- denom + crossprod(dZ, PkronR %*% dZ)
      } # for i
      if (length(denom) == 1) {
        denom <- try(1 / denom)
      } else {
        denom <- try(chol2inv(chol(denom)))
      }
      if (inherits(denom, "try-error")) {
        stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem in Z update.  denom is not invertible.\n Type MARSSinfo('denominv') for more info. \n", sep = "")
        stopped.with.errors <- TRUE
        break
      }
      MLEobj.iter$par$Z <- denom %*% numer

      # not sure this is needed; is guarding against R returning a vector
      if (!is.matrix(MLEobj.iter$par$Z)) MLEobj.iter$par$Z <- matrix(MLEobj.iter$par$Z, dim(d$Z)[2], 1)
      par1$Z <- parmat(MLEobj.iter, "Z", t = 1)$Z

      # Start~~~~~~~~~~~~Error checking
      if (control$safe) {
        new.kf <- rerun.kf("Z", MLEobj.iter, iter)
        if (!new.kf$ok) {
          stopped.with.errors <- TRUE
          msg.kf <- c(msg.kf, new.kf$msg.kf)
          stop.msg <- new.kf$stop.msg
          break
        } else {
          kf <- new.kf$kf
          MLEobj.iter$kf <- kf
          MLEobj.iter$logLik <- kf$logLik
          Ey <- MARSShatyt(MLEobj.iter)
          MLEobj.iter$Ey <- Ey
          msg.kem <- c(msg.kem, new.kf$msg.kem)
        }
      }
      if (control$trace > 0) {
        Ck <- kappa(denom)
        if (Ck > condition.limit) msg.kem <- c(msg.kem, paste("iter=", iter, " Unstable Z estimate because P_{t,t} is ill-conditioned. C =", round(Ck), sep = ""))
      }
    } # if !is.fixed Z
  } # end inner iter loop
  if (control$silent == 2) cat("\n")

  # prepare the MLEobj to return which has the elements set here
  MLEobj.return <- MLEobj
  MLEobj.return$iter.record <- iter.record
  MLEobj.return$numIter <- iter

  if (stopped.with.errors) {
    if (control$silent == 2) cat("Stopped due to numerical instability or errors. Print $errors from output for info or set silent=FALSE.\n")
    # print brief msg.  Full msg printed if silent=F
    msg <- c(stop.msg, "par, kf, states, iter, loglike are the last values before the error.\n")
    if (!control$safe) {
      msg <- c(msg, "Try control$safe=TRUE which uses a slower but slightly more robust algorithm.\n")
    }
    if (!control$trace > 0) {
      msg <- c(msg, "Use control$trace=1 to generate a more detailed error report. See user guide for insight.\n")
    }
    ## Attach any algorithm errors to the MLEobj
    if (control$trace > 0 && !is.null(msg.kem)) msg <- c(msg, "\nMARSSkem errors. Type MARSSinfo() for help.\n", msg.kem)
    if (control$trace > 0 && !is.null(msg.kf)) msg <- c(msg, "\nMARSSkf errors. Type MARSSinfo() for help.\n", msg.kf, "\n")
    MLEobj.return$errors <- msg

    MLEobj.return$par <- MLEobj.iter$par
    MLEobj.return$kf <- kf.last
    MLEobj.return$states <- kf.last$xtT
    MLEobj.return$convergence <- 52
    MLEobj.return$logLik <- MLEobj.iter$logLik
    return(MLEobj.return)
  } # stopped with errors

  ########### Did not stop with errors
  ## Set the convergence information
  ## Output depends on how it converged and how iterations were determined.  Possible conv.test$convergence values are
  #  MLEobj.iter$conv.test$convergence==72  means no convergense testing done at all because exited with errors BEFORE minit
  #  MLEobj.iter$conv.test$convergence==4 means stopped by hitting maxit; abstol not reached so no log-log testing done
  #  MLEobj.iter$conv.test$convergence==3 means abstol reached but no log-log test since control$min.iter.conv.test not reached
  #  MLEobj.iter$conv.test$convergence==0  means abstol reached and log-log test passed (=CONVERGED)
  #  MLEobj.iter$conv.test$convergence==1  means stopped by hitting maxit; abstol reached but NOT log-log
  #  MLEobj.iter$conv.test$convergence==-1 or -2  means log-log test returned errors
  catinfo <- !control$silent || control$silent == 2
  MLEobj.return$convergence <- 72 # debugging should be changed below
  if (MLEobj.iter$conv.test$convergence == 72) {
    MLEobj.return$convergence <- 52
    msg.conv <- MLEobj.iter$conv.test$messages
    if (catinfo) cat(paste("Error! EM algorithm exited at iter=", iter, " before minit reached.\n Minit was ", control$minit, ".\n", sep = ""))
  }
  if (MLEobj.iter$conv.test$convergence < 0) {
    MLEobj.return$convergence <- 62
    msg.conv <- MLEobj.iter$conv.test$messages
    if (catinfo) cat(paste("Error! EM algorithm exited due to errors reported by log-log test function.\n", sep = ""))
  }
  # if it returns 4, then abstol never reached before maxit hit.  Abstol must be hit before log-log
  # test run.  Thus we need to run the log-log test here IF we have enough iterations
  if (MLEobj.iter$conv.test$convergence == 4) {
    tmp.msg <- paste("Warning! Reached maxit before parameters converged. Maxit was ", control$maxit, ".\n", sep = "")
    if (iter >= control$min.iter.conv.test) {
      # loglog test has not been run because abstol never reached
      loglog.test <- loglog.conv.test(iter.record, iter, deltaT = control$conv.test.deltaT, tol = control$conv.test.slope.tol)
      MLEobj.iter$conv.test$messages <- loglog.test$messages
      if (loglog.test$convergence == 0) {
        MLEobj.return$convergence <- 11 # log-log passed but not abstol
        if (catinfo) cat(paste("Warning! log-log convergence only. Maxit (=", control$maxit, ") reached before abstol convergence.\n", sep = ""))
      }
      if (loglog.test$convergence < 0) {
        MLEobj.return$convergence <- 63
        if (catinfo) cat(paste(tmp.msg, " abstol not reached and log-log convergence returned errors.\n", sep = ""))
      }
      if (loglog.test$convergence == 1) {
        MLEobj.return$convergence <- 1
        if (catinfo) cat(paste(tmp.msg, " neither abstol nor log-log convergence tests were passed.\n", sep = ""))
      }
      if (loglog.test$convergence > 1) { # loglog test should never return > 1
        MLEobj.return$convergence <- 72
        if (catinfo) cat(paste(tmp.msg, " abstol not reached and log-log convergence returned errors.\n", sep = ""))
      }
    } else { # can'd do log-log test
      MLEobj.return$convergence <- 12 # can't do the log-log test
      if (catinfo) cat(paste(tmp.msg, " abstol not reached and no log-log test info since maxit less than min.iter.conv.test.\n", sep = ""))
    }
  }
  # At this point $conv.test$convergence==4 still and $convergence is 52,62,72 errors;
  # 1 neither abstol nor loglog; 11 loglog only; 12 no abstol and no info on loglog;

  # if conv.test$convergence == 3, abstol reached but no info on loglog since maxit < min.iter.conv.test
  if (MLEobj.iter$conv.test$convergence == 3) {
    tmp.msg <- paste("Warning! Abstol convergence only. no info on log-log convergence.\n", sep = "")
    MLEobj.return$convergence <- 3
    if (catinfo) cat(paste(tmp.msg, " Maxit (=", control$maxit, ") < min.iter.conv.test (=", control$min.iter.conv.test, ") so not log-log test.\n", sep = ""))
  }
  # if conv.tst$convergence==1, then abstol reached but no loglog convergence
  if (MLEobj.iter$conv.test$convergence == 1) {
    MLEobj.return$convergence <- 10
    msg.conv <- MLEobj.iter$conv.test$messages
    if (catinfo) cat(paste("Warning! Abstol convergence only. Maxit (=", control$maxit, ") reached before log-log convergence.\n", sep = ""))
  }
  if (MLEobj.iter$conv.test$convergence == 0) {
    MLEobj.return$convergence <- 0
    if (catinfo) {
      if (iter == control$minit) {
        cat(paste("Success! algorithm run for ", iter, " iterations. abstol and log-log tests passed.\n", sep = ""))
      } else {
        cat(paste("Success! abstol and log-log tests passed at ", iter, " iterations.\n", sep = ""))
      }
      if (control$conv.test.slope.tol > 0.1) cat(paste("Alert: conv.test.slope.tol is ", control$conv.test.slope.tol, ".\nTest with smaller values (<0.1) to ensure convergence.\n", sep = ""))
    } # !silent
  }
  msg.conv <- MLEobj.iter$conv.test$messages
  if (!is.null(msg.conv)) msg <- c(msg, "\nConvergence warnings\n", msg.conv)
  ##############################################################

  ## Other misc output
  MLEobj.return$par <- MLEobj.iter$par
  MLEobj.return$states <- MLEobj.iter$kf$xtT
  MLEobj.return$logLik <- MLEobj.iter$logLik

  if (!is.null(msg.kem)) {
    msg.kem <- c("\nMARSSkem warnings. Type MARSSinfo() for help.\n", msg.kem)
    msg <- c(msg, msg.kem)
  }
  if (!is.null(msg.kf)) {
    msg.kf <- c("\nMARSSkf warnings. Type MARSSinfo() for help.\n", msg.kf)
    msg <- c(msg, msg.kf)
  }
  if ((!is.null(msg.kem) || !is.null(msg.kf)) && control$trace < 1) {
    msg <- c(msg, "\nUse control$trace=1 to generate a more detailed error report.\n")
  }
  if ((!is.null(msg.kem) || !is.null(msg.kf)) && (!control$silent || control$silent == 2)) {
    cat("Alert: Numerical warnings were generated. Print the $errors element of output to see the warnings.\n")
  }

  ## Attach any algorithm errors to the MLEobj
  MLEobj.return$errors <- msg

  return(MLEobj.return)
}

## Run log-log convergence diagnostics
# 0 converged; 1 not converged; negative problem
loglog.conv.test <- function(iter.record, iter, params.to.test = c("Z", "U", "x0", "R", "Q", "A", "logLik"), deltaT = 9, tol = 0.5) {
  if (!is.list(iter.record) || !all(c("par", "logLik") %in% names(iter.record)) ||
    !any(params.to.test %in% c(names(iter.record$par), names(iter.record))) ||
    length(dim(iter.record$par)) != 2 || dim(iter.record$par)[1] <= 1 || is.null(colnames(iter.record$par))) {
    msg <- "par list not a proper list (with par and logLik) or too short for conv test or has no column names.\n"
    return(list(convergence = -1, messages = msg))
  } else {
    if ("logLik" %in% params.to.test) {
      # exp because we don't want the log of the log and subtract mean so exp(LL) doesn't = Inf
      # we are looking at slope so doesn't matter if we sub off the mean
      iter.record.par <- cbind(iter.record$par, logLik = exp(iter.record$logLik - mean(iter.record$logLik)))
      # iter.record.par = cbind(iter.record$par,logLik=iter.record$logLik)
    } else {
      iter.record.par <- iter.record$par
    }
    names.iter <- colnames(iter.record.par)
    names.sub <- strsplit(names.iter, "\\.")
    num.names <- length(names.sub)
    p.elems <- NULL
    for (j in 1:num.names) p.elems <- c(p.elems, names.sub[[j]][1])
    num.varcov <- sum(p.elems %in% params.to.test)
    test.conv <- rep(0, num.names)
    for (j in 1:num.names) {
      if (p.elems[j] %in% params.to.test) {
        test.len2 <- dim(iter.record.par)[1]
        test.len1 <- max(1, test.len2 - deltaT)
        test.len <- (iter - min(test.len2 - 1, deltaT)):iter
        test.par <- abs(iter.record.par[test.len1:test.len2, j])
        if (any(test.par == 0)) test.par <- test.par + 1
        # test.loglog=lm(log(test.par)~log(test.len))
        test.loglog <- (log(test.par[length(test.par)]) - log(test.par[1])) / (log(test.len[length(test.len)]) - log(test.len[1]))
        # test.conv[j]=test.loglog$coef[2]
        test.conv[j] <- test.loglog
      }
    }
  }
  if (any(is.na(test.conv))) {
    msg <- "The log-log degeneracy test produced NAs.\n"
    return(list(convergence = -2, messages = msg))
  }
  if (!is.null(test.conv) && !any(is.na(test.conv)) && any(abs(test.conv) > tol)) {
    msg <- paste("Warning: the ", names.iter[abs(test.conv) > tol], " parameter value has not converged.\n")
    msg <- c(msg, "Type MARSSinfo(\"convergence\") for more info on this warning.\n")
    return(list(convergence = 1, messages = msg, not.converged.params = names.iter[abs(test.conv) > tol], converged.params = names.iter[abs(test.conv) <= tol]))
  } else {
    return(list(convergence = 0, messages = NULL, not.converged.params = names.iter[abs(test.conv) > tol], converged.params = names.iter[abs(test.conv) <= tol]))
  } # 0 means converged successfully
}

rerun.kf <- function(elem, MLEobj, iter) { # Start~~~~~~~~Error checking
  if (iter == 1) cvg2 <- 1 + MLEobj$control$abstol
  msg.kem <- NULL
  msg.kf <- NULL
  loglike.old <- MLEobj$logLik
  kf <- MARSSkf(MLEobj)
  if (MLEobj$control$demean.states) {
    xbar <- apply(cbind(kf$x0T, kf$xtT), 1, mean)
    kf$xtT <- kf$xtT - xbar
    kf$x0T <- kf$x0T - xbar
  }
  if (!kf$ok) {
    msg.kf <- paste("iter=", iter, " ", elem, " update ", kf$errors, sep = "")
    stop.msg <- paste("Stopped at iter=", iter, " in MARSSkem after ", elem, " update: numerical errors in ", MLEobj$fun.kf, ".\n", sep = "")
    return(list(ok = FALSE, msg.kf = msg.kf, stop.msg = stop.msg))
  }
  loglike.new <- kf$logLik
  if (iter > 1 && is.finite(loglike.old) == TRUE && is.finite(loglike.new) == TRUE) cvg2 <- loglike.new - loglike.old
  if (iter > 2 & cvg2 < -sqrt(.Machine$double.eps)) {
    if (MLEobj$control$trace > 0) {
      msg.kem <- paste("iter=", iter, " LogLike DROPPED in ", elem, " update. logLik old=", loglike.old, " new=", loglike.new, "\n", sep = "")
    } else {
      msg.kem <- paste("MARSSkem: The soln became unstable and logLik DROPPED in the", elem, "updates.\n")
    }
  }
  return(list(kf = kf, msg.kem = msg.kem, msg.kf = msg.kf, ok = TRUE))
}

degen.test <- function(elem, MLEobj, iter) {
  if (is.fixed(MLEobj$marss$free[[elem]])) {
    return(list(MLEobj = MLEobj, msg = NULL, set.degen = FALSE))
  }
  # if elem is time-varying, don't allow degeneracy
  if (attr(MLEobj$model, "model.dims")[[elem]][3] != 1) {
    return(list(MLEobj = MLEobj, msg = NULL, set.degen = FALSE))
  }
  if (!MLEobj$control$allow.degen) {
    return(list(MLEobj = MLEobj, msg = NULL, set.degen = FALSE))
  }
  if (iter <= MLEobj$control$min.degen.iter) {
    return(list(MLEobj = MLEobj, msg = NULL, set.degen = FALSE))
  }
  if (!is.design(MLEobj$marss$free[[elem]], zero.rows.ok = TRUE)) {
    return(list(MLEobj = MLEobj, msg = NULL, set.degen = FALSE))
  } # strict, i.e. only 0 and 1
  isDiag <- function(x) {
    all(x[!diag(nrow(x))] == 0)
  }
  if (!isDiag(coef(MLEobj, type = "matrix", form = "marss")[[elem]])) {
    return(list(MLEobj = MLEobj, msg = NULL, set.degen = FALSE))
  } # only allow setting of 0s if diagonal

  # diagonal, not fixed, not time-varying, allow.degen set, iter>min iter and free is a design matrix
  # So can proceed
  # if here then not time-varying
  degen.par <- abs(MLEobj$par[[elem]]) < MLEobj$control$degen.lim
  msg.degen <- NULL
  set.degen <- FALSE
  dim.elem <- attr(MLEobj$marss, "model.dims")[[elem]][1]
  if (any(degen.par)) {
    for (i in which(degen.par)) {
      MLEobj.tmp <- MLEobj
      # set corresponding par to 0 and corresponding col of free to 0
      MLEobj.tmp$par[[elem]][i, 1] <- 0
      MLEobj.tmp$marss$free[[elem]][, i, 1] <- 0 # req not time-varying so set t=1
      # need to check that setting a R or Q diag to 0 doesn't lead to a improper model
      kemcheck <- MARSSkemcheck(MLEobj.tmp)
      if (kemcheck$ok) {
        new.kf <- MARSSkf(MLEobj.tmp)
        loglike.old <- MLEobj$logLik
        if (!new.kf$ok) msg.degen <- c(msg.degen, paste("iter=", iter, " MARSSkf returned error in attempt to set 0 diagonals for ", elem, "\n  ", new.kf$errors, "Perhaps Q and R are both going to 0?\n", sep = ""))
        if (new.kf$ok && is.finite(loglike.old) && is.finite(new.kf$logLik)) tmp.cvg2 <- new.kf$logLik - loglike.old else tmp.cvg2 <- Inf
        if (new.kf$ok && tmp.cvg2 < -sqrt(.Machine$double.eps)) {
          msg.degen <- c(msg.degen, paste("iter=", iter, " Setting diagonal to 0 blocked. logLik was lower in attempt to set 0 diagonals on ", elem, " logLik old=", loglike.old, " new=", new.kf$logLik, ", See MARSSinfo(\"", elem, "0blocked\").\n", sep = ""))
        }
        if (new.kf$ok && tmp.cvg2 > -sqrt(.Machine$double.eps)) { # this means degenerate elem has lower LL, so accept it
          MLEobj <- MLEobj.tmp
          MLEobj$kf <- new.kf
          MLEobj$Ey <- MARSShatyt(MLEobj) # needs the updated kf
          if (MLEobj$control$demean.states) {
            xbar <- apply(cbind(new.kf$x0T, new.kf$xtT), 1, mean)
            MLEobj$kf$xtT <- new.kf$xtT - xbar
            MLEobj$kf$x0T <- new.kf$x0T - xbar
          }
          MLEobj$logLik <- new.kf$logLik
          set.degen <- TRUE
        }
      } else {
        msg.degen <- c(msg.degen, paste("iter=", iter, " Setting element of ", elem, " to 0, blocked.  See MARSSinfo(\"", elem, "0blocked\"). The error is due to the following MARSSkemcheck errors.\n  MARSSkemcheck error: ", kemcheck$msg, sep = ""))
      }
    } # for degen.par; do one by one
  } # update MLEobj
  # set.degen is a flag to say if any 0 elements were set
  return(list(MLEobj = MLEobj, msg = msg.degen, set.degen = set.degen))
}


REBOL [
	Title: "REBOL Core extension"
	Purpose: {
		This extension contains commands not yet migrated to the core source code.
	}
	Author: "Richard Smolak"
	Copyright: {2010 2011 2012 2013 Saphirion AG, Zug, Switzerland}
	License: {
		Licensed under the Apache License, Version 2.0.
		See: http://www.apache.org/licenses/LICENSE-2.0
	}
	Name: core
	Type: extension
	Exports: [] ; added by make-host-ext.r
]

words: [
	n			;modulus
	e			;public exponent
	d			;private exponent
	p			;prime num 1
	q			;prime num 2
	dp			;CRT exponent 1
	dq			;CRT exponent 2
	qinv		;CRT coefficient
	priv-key	;private key
	pub-key		;public key
	g			;generator
	pkcs1		;padding type
]

init-words: command [
	words [block!]
]

init-words words

show-console: command [
	"Opens console window"
]

console-output: command [
	"Enables/Diables printing to console"
	state [logic!]
]

hide-console: command [
	"Hides console window if already open"
]

get-encap-data: command [
	"Returns encapped binary data"
]

to-png: command [
	"Save an image to PNG format"
	image [image!]
]

<no-export> req-dir: command [
	"low-level command used by REQUEST-DIR"
	/title
		text [string!]
	/path
		dir [string!]
]

request-dir: funct [
	"Asks user to select a directory and returns it as file path"
	/title
		"Custom dialog title text"
		text [string!]
	/path
		"Default directory path"
		dir [file!]
][
	if dir [dir: lib/replace/all to-local-file dir "/" "//"]
	if result: apply :req-dir [title text path dir] [
		return to-rebol-file result
	]
]

rc4: command [
	"Encrypt/decrypt data(modifies) using RC4 algorithm. Returns stream cipher context handle."
	/key
		crypt-key [binary!] "Crypt key. Have to be provided only for the first time to get stream handle!."
	/stream
		ctx [handle!] "Stream cipher context."
		data [binary! none!] "Data to encrypt/decrypt. Or NONE to close the cipher stream."
]

rsa-make-key: func [
	"Creates a key object for RSA algorithm."
][
	make object! [
		n:			;modulus
		e:			;public exponent
		d:			;private exponent
		p:			;prime num 1
		q:			;prime num 2
		dp:			;CRT exponent 1
		dq:			;CRT exponent 2
		qinv:		;CRT coefficient
		none
	]
]

rsa: command [
	"Encrypt/decrypt data using the RSA algorithm (with default PKCS#1 padding)."
	data [binary!]
	key-object [object!]
	/decrypt "Decrypts the data (default is to encrypt)"
	/private "Uses an RSA private key (default is a public key)"
	/padding "Selects the type of padding to use"
		padding-type [word! none!] "Type of padding. Available values: PKCS1 or NONE"
]

dh-make-key: func [
	"Creates a key object for Diffie-Hellman algorithm."
;NOT YET IMPLEMENTED
;	/generate
;		length [integer!] "Key length"
;	    generator [integer!] "Generator number"
][
	make object! [
		priv-key:	;private key
		pub-key:	;public key
		g:			;generator
		p:			;prime modulus
		none
	]
]

dh-generate-key: command [
    "Generates a new DH private/public key pair using given generator(g) and modulus(p) values."
    obj [object!] "The Diffie-Hellman key object"
]

dh-compute-key: command [
    "Computes the resulting, negotiated key from a private/public key pair and the peer's public key."
    obj [object!] "The Diffie-Hellman key object"
    public-key [binary!] "Peer's public key"
]

aes: command [
	"Encrypt/decrypt data using AES algorithm. Returns stream cipher context handle or encrypted/decrypted data."
	/key
		crypt-key [binary!] "Crypt key. Have to be provided only for the first time to get stream handle!."
		iv [binary! none!] "Optional initialization vector."
	/stream
		ctx [handle!] "Stream cipher context."
		data [binary! none!] "Data to encrypt/decrypt. Or NONE to close the cipher stream."
	/decrypt "Use the crypt-key for decryption (default is to encrypt)"
]


utilities_path = "./source/utilities.r"
source(utilities_path)
### Trajectory evaluation
#  The trajectory performance using monocle following these scripts(https://github.com/Winnie09/imputationBenchmark/blob/93f27e890a86fdc732257a4036bf38a52faf9f33/trajectory/code/hca/monocle2/01_get_score.R, https://github.com/Winnie09/imputationBenchmark/blob/93f27e890a86fdc732257a4036bf38a52faf9f33/trajectory/code/hca/tscan/01_get_score.R).
###  monocle2
library(monocle)
get_cds_monocle2 = function(gene_bc_mat){
  #  Make a new CDS and use DDRTree for dimension reduction.
  pd = new("AnnotatedDataFrame", data = data.frame(row.names = colnames(gene_bc_mat), cell = colnames(gene_bc_mat)))
  fd = new("AnnotatedDataFrame", data = data.frame(row.names = rownames(gene_bc_mat), gene_short_name = rownames(gene_bc_mat)))
  cds = newCellDataSet(gene_bc_mat, phenoData = pd, featureData = fd, expressionFamily = negbinomial.size())
  cds = estimateSizeFactors(cds)
  cds = estimateDispersions(cds)
  print("Reducing dimension...")
  cds = reduceDimension(cds)
  return(orderCells(cds))
}
get_score_monocle2 = function(cds, cell_type, correct_order, wrong_order, output_dir = NULL, type_level = NULL, show_interactively = FALSE){
  print("Looking for the root state...")
  used_cells = as.character(pData(cds)$cell)
  if(!is.null(output_dir) | show_interactively){
    if(!is.null(output_dir)){
      dir.create(output_dir, recursive = T, showWarnings = F)
    }
    p = plot_cell_trajectory(cds, color_by = "State")
    if(!is.null(output_dir)){
      ggsave(paste0(output_dir, "/state.pdf"), p)
    }
    if(show_interactively){
      print(p)
    }
    pData(cds)$CellType = cell_type[used_cells]
    p = plot_cell_trajectory(cds, color_by = "CellType")
    if(!is.null(output_dir)){
      ggsave(paste0(output_dir, "/celltype.pdf"), p)
    }
    if(show_interactively){
      print(p)
    }
    if(!is.null(type_level)){
      pData(cds)$Level = type_level[cell_type[used_cells]]
      p = plot_cell_trajectory(cds, color_by = "Level")
      if(!is.null(output_dir)){
        ggsave(paste0(output_dir, "/level.pdf"), p)
      }
      if(show_interactively){
        print(p)
      }
    }
  }
  cell_states = as.numeric(as.character(pData(cds)$State))
  names(cell_states) = used_cells
  unique_states = unique(cell_states)
  checkroot = sapply(unique_states, function(x){
    cds = orderCells(cds, root_state = x)
    return(length(cds@auxOrderingData[[cds@dim_reduce_type]]$root_cell))
  })
  candidate_root_states = sort(unique_states[checkroot > 0])
  result_mat = matrix(nrow = 0, ncol = 2, dimnames = list(c(), c("pair_number", "acc")))
  for(ii in candidate_root_states){
    cds = orderCells(cds, root_state = ii)
    this_output_dir = paste0(output_dir, "/rootstate_", ii)
    dir.create(this_output_dir, recursive = T, showWarnings = F)
    if(!is.null(output_dir) | show_interactively){
      p = plot_cell_trajectory(cds, color_by = "Pseudotime")
      if(!is.null(output_dir)){
        ggsave(paste0(this_output_dir, "/pseudotime.pdf"), p)
      }
      if(show_interactively){
        print(p)
      }
    }
    all_branch_points = cds@auxOrderingData[[cds@dim_reduce_type]]$branch_points
    if(length(all_branch_points) > 0){
      for(jj in seq(length(all_branch_points))){
        cds_tmp = cds
        tryCatch({
          cds_reduced = buildBranchCellDataSet(cds_tmp, branch_point = jj)
          df = data.frame(pData(cds_reduced),stringsAsFactors = F)[used_cells, ]
          if(!is.null(output_dir) | show_interactively){
            pData(cds_tmp)$Pseudotime = df[, "Pseudotime"]
            pData(cds_tmp)$Branch = df[, "Branch"]
            pData(cds_tmp)$State = df[, "State"]
            p = plot_cell_trajectory(cds_tmp, color_by = "Pseudotime")
            if(!is.null(output_dir)){
              ggsave(paste0(this_output_dir, "/branchpoint_", jj, "_pseudotime.pdf"), p)
            }
            if(show_interactively){
              print(p)
            }
            p = plot_cell_trajectory(cds_tmp, color_by = "Branch")
            if(!is.null(output_dir)){
              ggsave(paste0(this_output_dir, "/branchpoint_", jj, "_branch.pdf"), p)
            }
            if(show_interactively){
              print(p)
            }
            p = plot_cell_trajectory(cds_tmp, color_by = "State")
            if(!is.null(output_dir)){
              ggsave(paste0(this_output_dir, "/branchpoint_", jj, "_state.pdf"), p)
            }
            if(show_interactively){
              print(p)
            }
            pData(cds_tmp)$CellType = df[, "CellType"]
            p = plot_cell_trajectory(cds_tmp, color_by = "CellType")
            if(!is.null(output_dir)){
              ggsave(paste0(this_output_dir, "/branchpoint_", jj, "_celltype.pdf"), p)
            }
            if(show_interactively){
              print(p)
            }
            if(!is.null(type_level)){
              pData(cds_tmp)$Level = df[, "Level"]
              p = plot_cell_trajectory(cds_tmp, color_by = "Level")
              if(!is.null(output_dir)){
                ggsave(paste0(this_output_dir, "/branchpoint_", jj, "_level.pdf"), p)
              }
              if(show_interactively){
                print(p)
              }
            }
          }
          df = df[order(df$Pseudotime), ]
          score = rowSums(sapply(unique(df$Branch),function(x){
            branch_cell = as.character(df[df$Branch == x, 1])
            branch_celltype = cell_type[branch_cell]
            index_pair = combn(length(branch_cell), 2)
            if(min(index_pair[2,] - index_pair[1,]) < 0){
              stop("index_pair error")
            }
            branch_cellorder = sprintf("%s_%s",branch_celltype[index_pair[1, ]], branch_celltype[index_pair[2, ]])
            return(c(sum(branch_cellorder %in% correct_order), sum(branch_cellorder %in% wrong_order)))
          }))
          pair_number = sum(score)
          acc = score[1] / pair_number
          this_branch_point_results = matrix(c(pair_number, acc), nrow = 1)
          rownames(this_branch_point_results) = paste0("RS_", ii, "_BP_", jj)
          result_mat = rbind(result_mat, this_branch_point_results)
        }, error = function(e){
          cat(ii, " - ", jj, "\n")
          print(e)
        })
      }
    }
  }
  return(result_mat)
}

###
cell_level_df = data.frame(level = c(1, 2, 3, 3, 4, 4, 4, 5, 5, 5, 5, 5, 5),
                           immunepath = c(1, 1, 1, 0, 1, 0, 0, 1, 1, 1, 1, 0, 0),
                           monopath = c(1, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0),
                           erypath = c(1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1),
                           stringsAsFactors = F)
rownames(cell_level_df) = c("HSC", "MPP", "LMPP", "CMP", "CLP", "GMP", "MEP", "Bcell", "CD4Tcell", "CD8Tcell", "NKcell", "Mono", "Ery")
path_name = c("immunepath", "monopath", "erypath")
order_list = list(correct = list(), wrong = list(), cell_type = list())
for(ii in path_name){
  path_celltype = rownames(cell_level_df)[cell_level_df[, ii] == 1]
  order_list[["cell_type"]][[ii]] = path_celltype
  cell_type_pair = as.matrix(apply(expand.grid(path_celltype, path_celltype), 2, as.character))
  cell_type_pair = cell_type_pair[cell_type_pair[, 1] != cell_type_pair[, 2], ]
  correct_mat = cell_type_pair[cell_level_df[cell_type_pair[, 1], "level"] < cell_level_df[cell_type_pair[, 2], "level"], ]
  wrong_mat = cell_type_pair[cell_level_df[cell_type_pair[, 1], "level"] > cell_level_df[cell_type_pair[, 2], "level"], ]
  order_list[["correct"]][[ii]] = apply(correct_mat, 1, paste, collapse = "_")
  order_list[["wrong"]][[ii]] = apply(wrong_mat, 1, paste, collapse = "_")
}
type_level = as.character(cell_level_df[, "level"])
names(type_level) = rownames(cell_level_df)
correct_order_all = unique(unlist(order_list[["correct"]]))
wrong_order_all = unique(unlist(order_list[["wrong"]]))


########### R script to analyze historic weather data for temperature
## Connection handle con to BLU for Cloud data warehouse is provided already
## For plotting, we are using ggplot2 package
##
## Data for multiple stations is shown in different colors
##
library(ggplot2)
library(ibmdbR)

## initialize DB2 connection and environment
con <- idaConnect("BLUDB","","")
idaInit(con)

## query DB2 weather data and fetch min/max values of min/max values
## (lower/upper boundary each)
query<-paste('select max(lufttemperatur_maximum) as maxmax,min(lufttemperatur_minimum) as minmin,min(lufttemperatur_maximum) as minmax,max(lufttemperatur_minimum) as maxmin,tag,sid from (select lufttemperatur_maximum, lufttemperatur_minimum, day(mdatum) as tag,sid from blu06039.klima where month(mdatum)=9) group by tag,sid order by tag asc')
df <- idaQuery(query,as.is=F)

## Some plotting needs to be done
jpeg(type='cairo',"tempe.jpg",width=800,height=600)
ggplot(df, aes(x = TAG))+ylab("Temperature")+xlab("Day")+
     geom_line(aes(y = MAXMAX, group=SID, colour=SID))

sink('/dev/null')

idaClose(con)
## connection is closed, we are done


require(rms)
require(xlsx)
require(data.table)
require(Hmisc)

tablename<-outFile
ParentDataFolders<-c(parFile1, strsplit(parFile3, ","))
setwd(parFile1)

ctrFile<-parFile2

load(ctrFile)
options("stringsAsFactors"=FALSE)
out1=OutputSummary(ParentDataFolders,ctr.opo)

GenerateTables(out1,csv=T,tablename=tablename)

getci=function(x)
{
  x=as.numeric(x)[2:11]
  paste0(round(mean(x),1)," (",round(mean(x)-1.96*sd(x),1),", ",round(mean(x)+1.96*sd(x),1),")")
}

T2<-NULL
for (idx in c(1:length(out1))){
  T2.idx=data.frame(LocId=out1[[idx]]$Tx.Total.DSA$Table.Tx.DSA$PatLocId,CI=apply(out1[[idx]]$Tx.Total.DSA$Table.Tx.DSA,1,getci))
  colnames(T2.idx)<-c("LocId", attr(out1[[idx]], "simName"))
  if(is.null(T2)){
    T2<-T2.idx
  }else{
    T2<-merge(T2,T2.idx,by="LocId")
  }
}

write.csv(T2,file=paste0(tablename, "_transplantTo.csv"))


library(shiny)
library(shinyjs)
library(shinydashboard)
library(data.table)
library(DT)
library(DCGen)

source("src/elicitationIO.r")

server <- function(input, output, session) {
  vals <- reactiveValues()

  vals$Data <- getDefaultInputData()
  vals$LastSelected <- 1
  vals$NextID <- 4

  vals$Experts <- NULL
  vals$Pars <- NULL
  vals$Curves <- NULL

  disable("downloadReport")
  disable("downloadParameters")
  disable("downloadCurves")

  output$ElicitationBox <- renderUI({
    box(
      width = 12, collapsible = TRUE, title = "Elicitation Data",
      column(
        6,
        offset = 3,
        HTML('
'),
        actionButton(inputId = "Add_expert", label = "Add a new expert"),
        HTML('
')
      ),

      column(12, dataTableOutput("ElicitationData")),
      tags$script(
        HTML(
          '$(document).on("click", "input", function () {
          var checkboxes = document.getElementsByName("expert_selected");
          var checkboxesChecked = checkboxes.filter(function(d) {return d.checked;}).map(function(d) {return d.value});

          Shiny.onInputChange("checked_rows",checkboxesChecked);
  })'
        )
      ),
      tags$script(
        "$(document).on('click', '#ElicitationBox button', function () {
        Shiny.onInputChange('lastClickId',this.id);
        Shiny.onInputChange('lastClick', Math.random())
});"
      )

    )
  })

  output$ElicitationData <- renderDataTable({
    DT = showInputData(vals$Data)
    DT[["Actions"]] <-
      paste0(
        '

        <button type="button" class="btn btn-secondary modify"id=modify_', 1:nrow(vals$Data), '>Modify</button>
        <button type="button" class="btn btn-warning delete" id=delete_', 1:nrow(vals$Data), '>Delete</button>

        '
      )
    datatable(DT, escape = F)
  })

  modal_modify <- modalDialog(uiOutput("ModifyElicitationData"),
                              footer = tagList(actionButton("save_changes", "Save changes"), modalButton("Dismiss")),
                              size = "l")

  output$ModifyElicitationData <- renderUI({
    dat <- data.frame(vals$CurrentExpert)

    fluidRow(
      h3(strong("Set elicitation data"), align = "center"),
      column(12, textInput("new_name", "Expert name", value=dat[, 'Expert'], width="100%")),
      h4("Cap of adoption number", align = "center"),
      distribution_setter(dat[, 1:4+1], "M"),
      h4("Number of adoption in the first period", align = "center"),
      distribution_setter(dat[, 1:4+5], "N1"),
      h4("When the new adoption start to decrease", align = "center"),
      distribution_setter(dat[, 1:4+9], "T")
    )
  })

  observeEvent(input$lastClick,
               {
                 if (input$lastClickId %like% "delete") {
                   row_to_del = as.numeric(gsub("delete_", "", input$lastClickId))
                   vals$Data = vals$Data[-row_to_del]
                 } else if (input$lastClickId %like% "modify") {
                   vals$LastSelected <- as.numeric(gsub("modify_", "", input$lastClickId))
                   vals$CurrentExpert <- vals$Data[vals$LastSelected]
                   showModal(modal_modify)
                 }
               })

  observeEvent(input$Add_expert, {
    id <- vals$NextID
    vals$NextID <- id + 1
    new_row <- getNewInputData(id)
    vals$Data <- rbind(vals$Data, new_row)

    vals$LastSelected <- nrow(vals$Data)
    vals$CurrentExpert <- vals$Data[vals$LastSelected]

    showModal(modal_modify)
  })

  observeEvent(input$save_changes, {
    changed <- updateDistibution(input)
    names(changed) <- names(vals$Data)

    vals$Data[vals$LastSelected] <- changed
    vals$CurrentExpert <- vals$Data[vals$LastSelected]
    removeModal()
  })

  observeEvent({input$updateBass; input$dcg}, {
    vals$Experts <- inputELCs(vals$Data)

    nt <- input$num_iter
    # Create a Progress object
    progress <- shiny::Progress$new(session, min=1, max=nt)
    progress$set(message = "Translating Data", value = 0)

    on.exit(progress$close())

    updateProgress <- function(value=NULL, detail =NULL) {
      if (is.null(value)) {
        value <- progress$getValue()
        value <- value + (progress$getMax() - value) / 5
      }
      progress$set(value = value, detail = detail)
    }

    vals$Pars <- rand_parameters(vals$Experts, nt, method=input$agg_method, type="continuous")

    vals$Curves <- generate_diffusion_curves(vals$Pars, t_max=input$projT, progress=updateProgress, dt=0.5)

    vals$CurrentT <- input$projT

    shinyjs::enable("downloadReport")
    shinyjs::enable("downloadParameters")
    shinyjs::enable("downloadCurves")
  })

  output$PlotDistNT <- renderPlot({
    if (is.null(vals$Pars)) return ()
    visualise_inputs(vals$Pars)
  })

  output$PlotDistPQ <- renderPlot({
    if (is.null(vals$Pars)) return ()
    visualise_fitted(vals$Pars)
  })

  output$PlotBass <- renderPlot({
    curves <- vals$Curves
    if (is.null(curves)) return ()

    t_max <- max(ceiling(input$projT), 1)

    if (t_max > vals$CurrentT) {
      vals$Curves <- generate_diffusion_curves(vals$Pars, t_max=t_max, progress=updateProgress, dt=0.5)
    }

    visualise_curves(curves, dN=input$dnShow,  ci_range=input$centVal, statistics=input$curveType=="stats", average=input$avgType)

  })

  output$SummaryBass <- renderPrint({
    cat("Input parameters\n")
    print(summary(vals$Pars$Parameters[, c("M", "N1", "t")]))

    cat("\nOutput parameters\n")
    print(summary(vals$Pars$Parameters[, c("M", "p", "q")]))
  })

  output$downloadReport <- downloadHandler(
    filename = function() {
      paste(input$fileName, sep = ".", switch(
        input$format, PDF = "pdf", HTML = "html", Word = "docx"
      ))
    },

    content = function(file) {

      mcvs <- apply(vals$Curves, c(1, 2), ifelse(input$avgType=="mean", mean, median))
      mcvs <- round(mcvs, 2)
      mcvs <- mcvs[mcvs[, 1] == round(mcvs[, 1]),]
      mcvs <- mcvs[seq.int(1, nrow(mcvs), by=max(round(nrow(mcvs)/10), 1)), ]
      rownames(mcvs) <- rep("", nrow(mcvs))

      # src <- normalizePath("report_template.Rmd")
      # temporarily switch to the temp dir, in case you do not have write
      # permission to the current working directory
      # owd <- setwd(tempdir())
      # on.exit(setwd(owd))
      # file.copy(src, "report.Rmd", overwrite=TRUE)

      library(rmarkdown)
      format <-  switch(
        input$format,
        PDF=pdf_document(), HTML=html_document(), Word=word_document())

      #owd <- setwd(tempdir())
      #on.exit(setwd(owd))
      tempReport <- file.path(tempdir(), "report_template.Rmd")
      file.copy("report_template.Rmd", tempReport, overwrite = TRUE)

      out <- render(tempReport,
                    output_format=format,
                    intermediates_dir=tempdir(), knit_root_dir=tempdir(),
                    output_dir=tempdir(),
                    envir = list(dat=vals$Data, pars=vals$Pars, curves=vals$Curves, avg_curves=mcvs, input=input))

      file.copy(out, file)

    },
    contentType = "text/plain"
  )

  output$downloadParameters <- downloadHandler(
    filename = "DCG_Parameters.csv",

    content = function(file) {
      write.csv(vals$Pars$Parameters, file)
    }
  )

  output$downloadCurves <- downloadHandler(
    filename = "DCG_Curves.csv",

    content = function(file) {
      cvs <- data.frame(Time=vals$Curves[,1,1], Nt=vals$Curves[, 2, ])
      write.csv(cvs, file)
    }
  )
}


map.fisheries.data = function(p, outdir,  FUN, yrs, variable='effort',probs=c(0,0.975),log.variable=F,offset,theta=75,pres=10) {

  x = logbook.db( DS="logbook" )

  x = x [polygon_inside( x, region="isobath1000m"),]
  x = x[ which(x$effort <= 300) ,]
  x = x[ which(x$cpue < 600),]
  x$year=x$yr #this creates proper offset for 4X, 2017-18 season =2017
  x = lonlat2planar( x,  proj.type=p$aegis_proj4string_planar_km )
  x = subset(x,select=c('year','plon','plat',variable))
  names(x)[4] = 'z'

  if(variable=='landings')x$z = x$z/1000

  predlocs = bathymetry.db(p=p, DS="baseline")
  er = quantile( x$z[x$z>0], probs=probs)  # range of all years
  ler = er

  if(missing(offset))offset = er[1]  # offset fot log transformation
  if(missing(yrs))yrs=sort(unique(x$year))

  xyz = list()

  for(i in 1:length(yrs)){

    xd = subset(x,year==yrs[i],c('plon','plat','z'))

    p$pres=pres
    g= spatial_grid(p,"planar.coords")
    xd$plon = grid_internal( xd$plon, g$plon )
    xd$plat = grid_internal( xd$plat, g$plat )
    xyz[[i]] = aggregate(z~plon+plat,data=xd,FUN)
    names( xyz[[i]]) = c("plon", "plat", "z")
  }

  # calculate data range to show
  ler=quantile( do.call("rbind",xyz)$z, probs=probs)
  datarange = seq( ler[1], ler[2], length.out=50)
  #ckey=list(labels=list(at=pretty(datarange,8),labels=pretty(datarange,8),cex=2))

  cols = colorRampPalette(c( "yellow", "orange","darkred", "black"), space = "Lab")
  #cols = colorRampPalette(c("darkblue","cyan","green", "yellow", "orange","darkred", "black"), space = "Lab")
  for(i in 1:length(yrs)){
    xyz[[i]]$z[xyz[[i]]$z>ler[2]] = ler[2] # set levels higher than max datarange to max datarange
    outfn = paste( variable,yrs[i], sep=".")
    outloc = file.path( outdir,variable)

    # do the map
    dir.create (outloc, showWarnings=FALSE, recursive =TRUE)
    fn = file.path( outloc, paste(outfn, "png", sep="." ) )
    png( filename=fn, width=3072, height=2304, pointsize=40, res=300 )
      lp = aegis_map( xyz[[i]], xyz.coords="planar", depthcontours=TRUE,
        pts=NULL, annot=yrs[i], annot.cex=4, at=datarange, colorkey=NULL,
        col.regions=cols(length(datarange)+1), colpts=FALSE, corners=p$corners, display=FALSE, rez=c(pres,pres), plotlines="cfa.regions")
    print(lp)
    dev.off()
    print(fn)

    #try( ) ##BZ Nov 2018- Removed. Producing errors. No expression to "try" within the try() function
  }
}


# Runs ExomeDepth over the datasets cofigured at [datasets_params_file]
#USAGE: Rscript runExomeDepth.r [ExomeDepth_params_file] [datasets_params_file]
print(paste("Starting at", startTime <- Sys.time()))
suppressPackageStartupMessages(library(yaml))
suppressPackageStartupMessages(library(ExomeDepth))
library(methods)
source(if (basename(getwd()) == "optimizers") "../utils/utils.r" else "utils/utils.r") # Load utils functions

# returns gene for position matching ROI in bed data
auxGetGene <- function(bedData, chr, pos){
  gene <- bedData[bedData$chr == chr & (pos >= bedData$start & pos <= bedData$end),"gene"]
  if (length(gene) == 0)
    return(NULL)
  else
    return(gene)
}

# Process ExomeDepth alg. body
processExomedepthBody <- function(testCountsDF, controlCountsDF, countsDef, params){
  all <- data.frame()

  for (i in 1:ncol(testCountsDF)) {

    # build necessary matrix
    sampleName = colnames(testCountsDF)[i]
    if (sampleName %in% colnames(controlCountsDF)){
      controlCounts <- as.matrix(controlCountsDF)[, -i]
    } else
      controlCounts <- as.matrix(controlCountsDF)
    testCounts <- as.matrix(testCountsDF)[,i]

    # define final controls set
    references = select.reference.set(test.counts = testCounts,
                                      reference.count = controlCounts,
                                      bin.length = (countsDF$end - countsDF$start) / 1000,
                                      n.bins.reduced = 10000,
                                      phi.bins = params$phi.bins)
    controls = apply(X = as.matrix(controlCounts[, references$reference.choice]), MAR=1, FUN=sum)

    # call cnvs
    all_exons = new('ExomeDepth',
                    test = testCounts,
                    reference = controls,
                    formula = 'cbind(test, reference) ~ 1')
    all_exons = CallCNVs(x = all_exons,
                         transition.probability = params$transition.probability,
                         expected.CNV.length = params$expected.CNV.length,
                         chromosome = countsDef$space,
                         start = countsDef$start,
                         end = countsDef$end,
                         name = countsDef$names)

    if (nrow(all_exons@CNV.calls) > 0){
      # add sample column: remove first char "X"
      all_exons@CNV.calls$sample <- sampleName

      # add gene column
      for(i in 1:nrow(all_exons@CNV.calls)) {
        row <- all_exons@CNV.calls[i,]
        row$gene <- auxGetGene(bedData, row$chromosome, row$end)
        if (is.null(row$gene))
          stop(paste("Error: gene not found for chromosome", row$chromosome, ", start", row$start, ", end", row$end))

        all_exons@CNV.calls[i, "Gene"] <- row$gene
      }
    }

    all <- rbind(all, all_exons@CNV.calls)
  }

  return(all)
}

# Read args
args <- commandArgs(TRUE)
print(args)
if(length(args)>0) {
  exomedepthParamsFile <- args[1]
  datasetsParamsFile <- args[2]
} else {
  exomedepthParamsFile <- "algorithms/exomedepth/exomedepthParams.yaml"
  datasetsParamsFile <- "datasets.yaml"
}

#Load the parameters file
params <- yaml.load_file(exomedepthParamsFile)
datasets <- yaml.load_file(datasetsParamsFile)

# print params
print(paste("Params for this execution:", list(params)))

# go over datasets and run ExomeDepth for those which are active
for (name in names(datasets)) {

  dataset <- datasets[[name]]

  if (dataset$include){
    print(paste("Starting ExomeDepth for", name, "dataset", sep=" "))

    # extract fields
    bamsDir <- file.path(dataset$bams_dir)
    bedFile <- file.path(dataset$bed_file)
    bedData <- read.table(bedFile, sep="\t", stringsAsFactors=FALSE, col.names = (c("chr", "start", "end", "gene")))
    fastaFile <- file.path(dataset$fasta_file)

    # set readlength from algorithm params if defined
    if (!is.null(params$readLength)){
      readLength <- params$readLength
    } else
      readLength <- dataset$read_length

    # build output folder and file
    if (!is.null(params$outputFolder)) {
      outputFolder <- params$outputFolder
    } else
      outputFolder <- file.path(getwd(), "output", paste0("exomedepth-", name))
    if (is.null(params$execution) || params$execution != "skipPrecalcPhase") {
      unlink(outputFolder, recursive = TRUE);
      dir.create(outputFolder)
    }
    outputFile <- file.path(outputFolder, "all_cnv_calls.txt")
    failedRegionsFile <- file.path(outputFolder, "failedROIs.csv")

    # Do pre-calc part of the algorithm
    if (is.null(params$execution) || params$execution != "skipPrecalcPhase") {

      # read bam counts
      bamFiles <- list.files(bamsDir, "*.bam$", full.names=T)
      counts <- getBamCounts(bed.file = bedFile,
                             bam.files = bamFiles,
                             read.width = readLength,
                             referenceFasta = fastaFile)
      countsDF  <- as.data.frame(counts)
      names(countsDF) <- gsub(".bam", "", names(countsDF)) # remove .bam from sample name

      # Fix rare bug: sometimes X is added at the beginning of input samples
      parts <- strsplit(bamFiles[1], "/")
      original <- gsub(".bam", "", parts[[1]][length(parts[[1]])])
      if (substr(original, 1, 1) != "X" && substr(names(countsDF)[7], 1, 1) == "X")
        names(countsDF)[7:ncol(countsDF)] <- substring(names(countsDF)[7:ncol(countsDF)] , 2)

      if (!is.null(params$execution) && params$execution == "onlyPrecalcPhase") {
        # save results to be used by other executions
        dir.create(params$precalcFolder)
        saveRDS(countsDF, file.path(params$precalcFolder, "countsDF.rds"))

        print(paste("ExomeDepth (Only pre-calc phase) for", name, "dataset finished", sep=" "))
        cat("\n\n\n")
        quit()
      }
    } else {  # skipPrecalcPhase mode: read previous results
      print(paste("ExomeDepth Skipping pre-calc phase for", name, "dataset finished", sep=" "))
      countsDF <- readRDS(file.path(params$precalcFolder, "countsDF.rds"))
    }

    # Process each sample
    all <- data.frame()
    if (!is.null(dataset$clinical_indication) && dataset$clinical_indication != ""){
      sample_indications <- read.table(dataset$clinical_indication, header = T, sep = "\t", stringsAsFactors=F)[,c(1,2)]
      sample_indications <- sample_indications[sample_indications$Genes != "" & sample_indications$Genes != " ", ] # remove empty items
      indicationsComp <- readCompatibleIndications(dataset$clinical_indication) # get compatible indications for each indication

      for (indication in names(indicationsComp)) {
        print(paste("\nProcessing indication:", indication))

        # samples to test: samples matching this indication
        testSamplesNames <- sample_indications[sample_indications$Genes == indication,][["SampleID"]]
        if (length(testSamplesNames > 0)){

          # Calculate available control samples: those compatible with test indication
          compatibles <- indicationsComp[[indication]]
          samplesToExclude <- sample_indications[!sample_indications$Genes %in% compatibles,][["SampleID"]]
          controlSamplesNames <- sample_indications[!sample_indications$SampleID %in% samplesToExclude, ][["SampleID"]]

          # Call other algorithm steps
          results <- processExomedepthBody(testCountsDF = countsDF[testSamplesNames],
                                           controlCountsDF = countsDF[controlSamplesNames],
                                           countsDef = countsDF[1:6],
                                           params = params)
          all <- rbind(all, results)
        }
      }

    } else {
      message("Clinical indication not found, using all samples potentially as controls")

      all <- processExomedepthBody(testCountsDF = countsDF[7:ncol(countsDF)],
                                   controlCountsDF = countsDF[7:ncol(countsDF)],
                                   countsDef = countsDF[1:6],
                                   params = params)
    }

    # write results
    write.table(all, file = outputFile, sep='\t', row.names=FALSE, col.names=TRUE, quote=FALSE)

    df <- read.table(outputFile,sep="\t", stringsAsFactors=FALSE, header = TRUE) # REMOVE

    # Save results in GRanges format
    message("Saving GenomicRanges results")
    saveResultsFileToGR(outputFolder, "all_cnv_calls.txt", chrColumn = "chromosome", sampleColumn = "sample",
                        startColumn = "start", endColumn = "end", cnvTypeColumn = "type")

    # Save empty failed regions (not available in this algorithm)
    fr <- data.frame(matrix(ncol = 5, nrow = 0))
    colnames(fr) <- c("SampleID", "Chr", "Start", "End", "Gene")
    write.table(fr, failedRegionsFile, sep="\t", row.names=FALSE, quote = FALSE)

    print(paste("ExomeDepth for", name, "dataset finished", sep=" "))
    cat("\n\n\n")
  }
}

print(paste("Finishing at", endTime <- Sys.time()))
cat("\nElapsed time:")
print(endTime - startTime)


#NEON taxonomic group tables for all phylogenetic levels.
#common = found in greater than 50% of prior samples.
#clear environment, source paths, packages and functions.
rm(list=ls())
library(data.table)
source('paths.r')
source('NEFI_functions/common_group_quantification.r')

#set output path.----
output.path <- NEON_ITS_fastq_all_cosmo_phylo_groups_1k_rare.path

#load data.----
sv <- readRDS(NEON_ITS_fastq_SV.table_clean_1k_rare.path)
tax <- readRDS(NEON_ITS_fastq_fun_clean_1k_rare.path)

#get reference common phyla from prior.----
ref <- readRDS(tedersoo_ITS_common_phylo_groups_list_1k.rare.path)

#check that things are properly ordered.----
if(sum(rownames(tax) == colnames(sv)) != ncol(sv)){
  cat('Stop. rownames of tax-fun table do not match column names of SV table.\n')
}

#aggregate at multiple phylogenetic levels.----
levels <- c('phylum','class','order','family','genus','fg')
output <- list()
for(i in 1:length(levels)){
  level <- levels[i]
  reference <- data.table(ref[[i]]$group_frequencies)
  reference <- as.character(reference[sample_frequency > 0.5]$groups)
  output[[i]] <- common_group_quantification(sv, tax, reference, level, ref_filter = T)
}
names(output) <- names(ref)

#save output.----
saveRDS(output, output.path)


library(tidyverse)
# library(dplyr)
# library(ggrepel)
require(rgdal)
library(ggpubr)

#get folder directory
folder <- dirname(rstudioapi::getSourceEditorContext()$path)

results <- read.csv(file.path(folder, '..', '..', 'results', 'national_results_technology_options.csv'))
names(results)[names(results) == 'GID_0'] <- 'iso3'

clusters <- read.csv(file.path(folder, '..', 'clustering', 'results', 'data_clustering_results.csv'))
names(clusters)[names(clusters) == 'ISO_3digit'] <- 'iso3'
clusters <- select(clusters, iso3, cluster, country)

results <- merge(results, clusters, x.by='iso3', y.by='iso3', all=FALSE)

mean_user_cost_by_cluster <- select(results, scenario, strategy, confidence, cost_per_network_user, cluster)

mean_user_cost_by_cluster <- mean_user_cost_by_cluster %>%
  group_by(scenario, strategy, confidence, cluster) %>%
  summarise(mean_cost_per_user = mean(cost_per_network_user))

results <- merge(clusters, mean_user_cost_by_cluster, x.by='cluster', y.by='cluster', all=TRUE)

gdp <- read.csv(file.path(folder, '..', 'gdp.csv'))
names(gdp)[names(gdp) == 'iso3'] <- 'iso3'
gdp <- select(gdp, iso3, gdp)

results <- merge(results, gdp, by='iso3', all=FALSE)

pop <- read.csv(file.path(folder, 'data_inputs', 'population_2018.csv'))
pop <- select(pop, iso3, population)
pop$iso3 <- as.character(pop$iso3)

results <- merge(results, pop, by='iso3', all=FALSE)

results$total_cost <- results$mean_cost_per_user * results$population

results$gdp_percentage <- (results$total_cost / 5) / results$gdp * 100

results$confidence = factor(results$confidence, levels=c('5','50', '95'),
                             labels=c("lower", 'mean', "upper"))

results$scenario = factor(results$scenario, levels=c("S1_25_10_5",
                                                       "S2_200_50_25",
                                                       "S3_400_100_50"),
                           labels=c("S1 (25 Mbps)",
                                    "S2 (200 Mbps)",
                                    "S3 (400 Mbps)"))

# scenario_200_mbps <- results[(results$scenario == 'S2 (200 Mbps)'),]

results$strategy = factor(results$strategy, levels=c(
                                        "4G_epc_microwave_baseline_baseline_baseline_baseline",
                                        "4G_epc_fiber_baseline_baseline_baseline_baseline",
                                        "5G_nsa_microwave_baseline_baseline_baseline_baseline",
                                        "5G_sa_fiber_baseline_baseline_baseline_baseline"),
                                        labels=c(
                                          "4G (Microwave)",
                                          "4G (Fiber)",
                                          "5G NSA (Microwave)",
                                          "5G SA (Fiber)"))

results <- results[complete.cases(results), ]

wide <- select(results, country, cluster, scenario, strategy, confidence, gdp_percentage)
wide <- spread(wide, confidence, gdp_percentage)

gdp_perc_by_strategy <- ggplot(wide, aes(x=wide$cluster, y=wide$mean, colour=wide$cluster)) +
  geom_boxplot(aes(group=factor(wide$cluster))) +
  # geom_jitter(width = 0.4, height=0.5, size=1.7) +
  scale_colour_manual(values = c("#F0E442","#E69F00","#D55E00", "#0072B2", "#56B4E9","#009E73")) +
  theme(legend.position = NULL) +
  labs(colour=NULL,
       title = "Total Investment Cost",
       subtitle = "Results reported by scenario, strategy and country for 95% confidence",
       x = NULL, y = "Annual GDP (% over 5 years)") +
  theme(panel.spacing = unit(0.4, "lines")) +
  expand_limits(x = 0, y = 0) +
  scale_y_continuous(limits=c(0,10), expand=c(0,0), breaks=seq(0,9.7, 2)) +
  guides(fill=FALSE, colour=FALSE) +
  facet_grid(scenario~strategy)

# #export to folder
# path = file.path(folder, 'figures', 'gdp_perc_by_strategy.png')
# ggsave(path, units="cm", width=15, height=15)
# print(gdp_perc_by_strategy)
# dev.off()

if (!require(rgeos)) {
  install.packages("rgeos", repos = "http://cran.us.r-project.org")
  require(rgeos)
}
if (!require(rgdal)) {
  install.packages("rgdal", repos = "http://cran.us.r-project.org")
  require(rgdal)
}
# if (!require(raster)) {
#   install.packages("raster", repos = "http://cran.us.r-project.org")
#   require(raster)
# }
if(!require(ggplot2)) {
  install.packages("ggplot2", repos="http://cloud.r-project.org")
  require(ggplot2)
}
if(!require(viridis)) {
  install.packages("viridis", repos="http://cloud.r-project.org")
  require(viridis)
}
if(!require(dplyr)) {
  install.packages("dplyr", repos = "https://cloud.r-project.org/")
  require(dplyr)
}
if(!require(gtable)) {
  install.packages("gtable", repos = "https://cloud.r-project.org/")
  require(gtable)
}
if(!require(grid)) {
  install.packages("grid", repos = "https://cloud.r-project.org/")
  require(grid)
}
if(!require(readxl)) {
  install.packages("readxl", repos = "https://cloud.r-project.org/")
  require(readxl)
}
if(!require(magrittr)) {
  install.packages("magrittr", repos = "https://cloud.r-project.org/")
  require(magrittr)
}

#set ggplot2 theme for map
theme_map <- function(...) {
  theme_minimal() +
    theme(
      text = element_text(family = "Ubuntu Regular", color = "#22211d"),
      axis.line = element_blank(),
      axis.text.x = element_blank(),
      axis.text.y = element_blank(),
      axis.ticks = element_blank(),
      axis.title.x = element_blank(),
      axis.title.y = element_blank(),
      # panel.grid.minor = element_line(color = "#ebebe5", size = 0.2),
      panel.grid.major = element_line(color = "#ebebe5", size = 0.2),
      panel.grid.minor = element_blank(),
      plot.background = element_rect(fill = "#f5f5f2", color = NA),
      panel.background = element_rect(fill = "#f5f5f2", color = NA),
      legend.background = element_rect(fill = "#f5f5f2", color = NA),
      panel.border = element_blank(),
      ...
    )
}

#get folder directory
folder <- dirname(rstudioapi::getSourceEditorContext()$path)

#get data
data <- read.csv(file.path(folder,'..','clustering','results','data_clustering_results.csv'), stringsAsFactors = F)

shapes <- file.path(folder, '..','clustering','data_inputs', "global_countries.shp")

#fortify the data AND keep trace of the commune code! (Takes ~2 minutes)
gde_15 <- readOGR(shapes, layer = "global_countries")

#fortify the data
map_data_fortified <- fortify(gde_15, region = "ISO_3digit") %>% mutate(id = id)

#now join the thematic data
map_data <- map_data_fortified %>% left_join(data, by = c("id" = "ISO_3digit"))

#swap na label for 'no data'
map_data$cluster[is.na(map_data$cluster)] <- "No Data"

#create map
cluster_map <- ggplot() +
  geom_polygon(data = map_data, aes(fill = factor(cluster),
       x = long, y = lat, group = group), colour='grey', size = 0.2) +
  coord_equal() +  theme_map() +
  labs(x = NULL, y = NULL,  fill = 'Cluster', title = "Global Countries by Cluster",
       subtitle = "Clustering based on GDP Per Capita, Population Density and 4G Coverage") +
  theme(legend.position = "bottom") +
  guides(fill=guide_legend(ncol=8)) +
  scale_fill_manual(values = c("#F0E442","#E69F00","#D55E00", "#0072B2", "#56B4E9","#009E73", "#999999", "#000000"))

# #export to folder
# path = file.path(folder, 'figures', 'map.png')
# ggsave(path, units="cm", width=15, height=15)
# print(cluster_map)
# dev.off()

rm(data, gde_15, map_data, map_data_fortified)

gdp_perc_by_strategy_vis <- ggarrange(gdp_perc_by_strategy, cluster_map, ncol = 1, nrow = 2, align = c("hv"))

#export to folder
path = file.path(folder, 'figures', 'gdp_perc_by_strategy_panel.png')
ggsave(path, units="cm", width=25, height=25)
print(gdp_perc_by_strategy_vis)
dev.off()