get_ip_fit <- function(static, dynamic, ...,
                    terms="scale", intercept="no", rmdata=TRUE, log=TRUE){
  if (intercept=="yes") intercept <- 1 else intercept <- 0
  terms <- as.character(terms)
  if (length(terms)==1){
    if (terms=="set1") terms <- c("facility", "is_small", "season", "CaseType")
  }
  sel_sta <- c("outbreak_code", terms)
  sel_dyn <- c("outbreak_code", "symp_obs")
  mg <- merge(dynamic[, sel_dyn], static[, sel_sta], by="outbreak_code")
  rhs <- paste(c(intercept, terms), collapse="+")
  model <- paste("symp_obs", rhs, sep="~")
  ret <- list()
  # In case some observations are 0
  mg$symp_obs <- mg$symp_obs + smidge
  if (!log){
    ret$m <- try(glm(as.formula(model), family=Gamma(identity), data=mg))
  } else {
    ret$m <- try(glm(as.formula(model), family=Gamma(log), data=mg))
  }
  if (!inherits(ret$m, "try-error")){
    s <- summary(ret$m)
    ret$disp <- s$disp
    ret$coefs <- s$coefficients
    gs <- try(MASS::gamma.shape(ret$m))
    if (!inherits(gs, "try-error")){
      ret$disp_ml <- 1/gs$alpha
      ret$disp_ml_se <- gs$SE/gs$alpha^2
      ret$coefs_ml <- summary(ret$m, dispersion=ret$disp_ml)$coefficients
    }
    if (rmdata) {
      ret$m$data <- NULL
    }
  }
  ret
}

get_hazard_funs <- function(piece){
  combine_jumps <- function(foo){
    foo <- foo[order(foo[,"t"]), ]
    delta <- tapply(foo[, "delta"], foo[, "t"], sum)
    t <- unique(foo[, "t"])
    cbind(t, delta)
  }
  ret <- list()
  Ymat <- rbind(cbind(piece$infective_start, 1),
                cbind(piece$recovered_start, -1))
  colnames(Ymat) <- c("t", "delta")
  Ymat <- combine_jumps(Ymat)
  Y <- cumsum(Ymat[, "delta"])
  Y <- Y[-length(Y)]
  holding_times <- diff(Ymat[, "t"])
  Yt <- Y * holding_times
  jump_times <- Ymat[, "t"]
  ret$min_jump_time <- min(jump_times)
  ret$cum_hazard <- function(time){
    test <- piece$recovered_start < time
    end <- ifelse(test, piece$recovered_start, time)
    dif <- end - piece$infective_start
    if (any(dif > 0)){
      ret <- sum(dif[dif >0])
      ret <- unname(ret)
    } else {
      ret <- 0
    }
    ret
  }
  ret$instant_hazard <- function(time){
    test <- jump_times <= time
    if (any(test)){
      i <- sum(test)
      ret <- unname(Y[i])
    } else {
      ret <- 0
    }
    ret 
  }
  ret
}

get_est <- function(job, static, dynamic, ...,
                    gr_terms="scale", gr_intercept="no",
                    tr_terms=gr_terms, tr_intercept=gr_intercept,
                    ip_terms=gr_terms, ip_intercept=gr_intercept,
                    est_latent_period=1, method="admb", logip=TRUE,
                    limit_staff_to_one_day=TRUE, altfits=FALSE){
  ret <- list()
  types <- levels(dynamic$CaseType)
  ret$propFreeImp <- static$propFreeImp
  ret$propOnsetNA <- static$propOnsetNA

  ret$ip_fit <- get_ip_fit(static=static, dynamic=dynamic,
                           terms=ip_terms, intercept=ip_intercept, log=logip)
  
  dynamic$est_latent_start <- dynamic$infective_start - est_latent_period
  tmpf <- function(piece){
    data.frame(tau=sum(piece$recovered_start - piece$infective_start),
               OutbreakNumber=piece$OutbreakNumber[1])
  }

  if (limit_staff_to_one_day){
    one_day <- dynamic$infective_start + 1
    one_day <- ifelse(one_day < dynamic$recovered_start, one_day,
                      dynamic$recovered_start)
    test <- dynamic$CaseType == "staff"
    dynamic$recovered_start[test] <- one_day[test]
    rm(one_day, test)
  }
    
  dat <- ddply(dynamic, "outbreak_id", tmpf)
  hfuns <- dlply(dynamic, "outbreak_id", get_hazard_funs)
  dyns <- dlply(dynamic, "outbreak_id", identity)
  tmpf <- function(x, y){
    tmpff <- function(type){
      ret <- list()
      z <- y[y$CaseType==type, ]
      if (nrow(z) > 0){
        is_seed <- z$infective_start == x$min_jump_time
        ret$hij <- sapply(z$est_latent_start, x$cum_hazard)[!is_seed]
        ret$Yij <- sapply(z$est_latent_start, x$instant_hazard)[!is_seed]
      } else {
        ret$hij <- ret$Yij <- numeric(0)        
      }
      ret
    }
    ret <- lapply(types, tmpff)
    names(ret) <- types
    ret
  }
  vdat <- mapply(tmpf, x=hfuns, y=dyns, SIMPLIFY=FALSE)
  tmpf <- function(x){
    lapply(x, "[[", "Yij")
  }
  vYij <- lapply(vdat, tmpf)
  tmpf <- function(x) {
    lapply(x, function(y) data.frame(sumh=sum(y$hij), k=length(y$hij)))
  }
  foo <- lapply(vdat, tmpf)
  tmpf <- function(x, name) {
    ret <- ldply(x, identity)
    ret[, name] <- ret$.id
    ret$.id <- NULL
    ret
  }
  foo <- lapply(foo, tmpf, name="CaseType")
  foo <- tmpf(foo, name="outbreak_id")
  dat <- merge(foo, dat, by="outbreak_id")

  if(sum(dat$k) == 0) {
    fit <- list(solution_type="no_transmission_data")
    ret$fit <- list(admb=fit, uniroot=fit)
  } else if(sum(dat$sumh) == 0) {
    fit <- list(solution_type="exposure_rounded_to_zero")
    ret$fit <- list(admb=fit, uniroot=fit)
  } else {
    dat$outbreak_code <- with(dat, paste(OutbreakNumber, CaseType, sep="_"))
    get_mmat <- function(intercept, terms){
      if (intercept=="yes") intercept <- 1 else intercept <- 0
      terms <- as.character(terms)
      if (length(terms) == 1){
        if (terms=="set1") terms <- c("facility", "is_small",
              "season", "CaseType")
      }
      ord <- match(dat$outbreak_code, static$outbreak_code)   
      rhs <- paste(c(intercept, terms), collapse="+")
      form <- as.formula(paste("~", rhs))
      data <- static[ord, ]
      rownames(data) <- seq_len(nrow(data))
      model.matrix(form, data=data)
    }
    mmat_gr <- get_mmat(gr_intercept, gr_terms)
    mmat_tr <- get_mmat(tr_intercept, tr_terms)
    ret$coefnames <- list(mmat_tr=colnames(mmat_tr),
                         mmat_gr=colnames(mmat_gr))
    valid_rows <- intersect(rownames(mmat_tr), rownames(mmat_gr))
    dat <- dat[rownames(dat) %in% valid_rows, ]

    tmpf <- function(x, y){
      vYij[[x]][[y]]
    }
    datYij <- mapply(tmpf, x=dat$outbreak_id, y=dat$CaseType, SIMPLIFY=FALSE)
    stopifnot(all(sapply(datYij, length) == dat$k))

    ## To get minimum number of susceptibles in each group.
    ##
    ##   interact() could be used instread of pasting string codes,
    ##   but I seem to remember it being very slow
    codes1 <- apply(mmat_tr, 1, function(x) paste(x, collapse="_"))
    codes2 <- apply(mmat_gr, 1, function(x) paste(x, collapse="_"))
    codes <- paste(codes1, codes2, sep="_")
    dat$maxk <- NA
    split(dat, codes) <- lapply(split(dat, codes),
                                function(x) {x$maxk <- max(x$k); x})
    rm(codes, codes1, codes2)
    
    if (method=="admb") {
      prefix <- paste("admb_", job$id, "_", sep="")
      ret$fit <- list(admb=get_admb_fit(dat, datYij, mmat_gr, mmat_tr,
                        workdir=tempfile(prefix)))
      if (altfits==TRUE){
        prefix <- paste("admbcan_", job$id, "_", sep="")
        ret$fitcan <- list(admb=get_admb_fit(dat, datYij, mmat_gr, mmat_tr,
                             workdir=tempfile(prefix), fn="epiglmcan"))      
        prefix <- paste("admb2_", job$id, "_", sep="")
        ret$fit2 <- list(admb=get_admb_fit2(dat, datYij, mmat_gr, mmat_tr,
                          workdir=tempfile(prefix)))
      }
    } else if(method=="uniroot") {
      stopifnot(all(ncol(mmat_gr)==1, ncol(mmat_gr)==1))
      ret$fit <- list(uniroot=get_uniroot_fit(dat))
    } else if(method=="both"){
      stopifnot(all(ncol(mmat_gr)==1, ncol(mmat_gr)==1))
      ret$fit <- list(admb=get_admb_fit(dat, datYij, mmat_gr, mmat_tr),
                      uniroot=get_uniroot_fit(dat))
    }
  }
  ret
}

get_uniroot_fit <- function(dat, max_X0_hat=1e5){
  n <- nrow(dat)
  n_transmissions <- dat$k
  Y_integral <- dat$tau
  inner_sum_H <- dat$sumh
  sum_Y_integral <- sum(Y_integral)

  ret <- list()
  sum_transmissions <- sum(n_transmissions)
  if(sum_transmissions == 0){
    ret$beta_hat <- NA
    ret$X0_hat <- NA
    ret$solution_type <- "no_transmission_data"
  } else {
    min_X0_hat <- max(n_transmissions)    
    if(min_X0_hat == 1){
      ## There are an infinite number of MLE for beta and mu
      ret$beta_hat <- NA
      ret$X0_hat <- NA
      ret$solution_type <- "unidentifiable"
    } else {           
      m <- function(X0){
        XY_ints <- inner_sum_H + Y_integral * (X0 - n_transmissions)
        sum_transmissions / sum(XY_ints)        
      }
      score_X0 <- function(X0_hat){
        t1 <- n * digamma(X0_hat + 1)
        t2 <- sum(digamma(X0_hat - n_transmissions + 1))
        t3 <- m(X0_hat) * sum_Y_integral
        -t1 + t2 + t3
      }
      if (score_X0(min_X0_hat) >= 0){
        ret$beta_hat <- m(min_X0_hat)
        ret$X0_hat <- min_X0_hat
        ret$solution_type <- "minimum_feasible_X0"
      } else {
        foo <- sum_transmissions * sum(Y_integral * n_transmissions - inner_sum_H)
        bar <- -sum_Y_integral * sum((n_transmissions - 1) * n_transmissions / 2)
        foobar <- foo + bar
        divergent <- foobar < 0
        if(divergent){
          ret$beta_hat <- 0
          ret$X0_hat <- Inf
          ret$solution_type <- "divergent"
        }
        else {
          ans <- try(uniroot(f=score_X0, interval=c(min_X0_hat, max_X0_hat)))
          while(inherits(ans, "try-error")){
            max_X0_hat <- max_X0_hat * 10
            ans <- try(uniroot(f=score_X0, interval=c(min_X0_hat, max_X0_hat)))
          }
          ret$X0_hat <- ans$root
          ret$beta_hat <- m(ret$X0_hat)
          ret$solution_type <- "interior_point"
        }
      }
    }
  }
  ret$coefficients <- c(ctr=ret$beta_hat, cgr=ret$X0_hat * ret$beta_hat)
  ret$loglik <- NA
  if (ret$solution_type == "interior_point"){
    fbeta <- sum(n_transmissions) / ret$beta_hat^2
    fX0 <- -n*trigamma(ret$X0_hat + 1)
    fX0 <- fX0 + sum(trigamma(ret$X0_hat - n_transmissions + 1))
    fbetaX0 <- sum_Y_integral
    ret$nllhessian <- rbind(c(fbeta, fbetaX0),
                            c(fbetaX0, fX0))

    Sigma <- try(solve22(ret$nllhessian))
    if(!inherits(Sigma, "try-error")){
      G <- rbind(c(1, 0),
                 c(ret$X0_hat, ret$beta_hat))
      ret$vcov <- G %*% Sigma %*% t(G)
      colnames(ret$vcov) <- rownames(ret$vcov) <- names(ret$coefficients)
      ret$se <- sqrt(diag(ret$vcov))

      foo <- cbind(ret$se, ret$se)
      ret$cor <- ret$vcov /(foo * t(foo))
    }
  }
  ret  
}

solve22 <- function(M){
  A <- M[1,1];  B <- M[1,2]
  C <- M[2,1];  D <- M[2,2]
  det <- A*D - B*C
  stopifnot(det!=0)
  rbind(c(D, -B),
        c(-C, A)) / det  
}

get_admb_fit <- function(dat, datYij, mmat_gr, mmat_tr,
                         workdir=tempfile("admb_dir_"),
                         fn="epiglm"){
  ## So simultaneous ADMB runs cannot clobber each other's files
  dir.create(workdir, recursive=TRUE)
  original_dir <- getwd()
  setwd(workdir)
  on.exit(setwd(original_dir))

  stopifnot(Sys.info()["sysname"] == "Linux")
  machext <- paste("-",   Sys.info()["machine"], sep="")
  pkg.root <- system.file(package="EpiGLM")
  lns <- function(suf){
    targ <- paste(fn, suf, sep="")
    src <- file.path(pkg.root, "admb", targ)
    file.copy(src, targ)    
  }
  sufs <- c(machext, ".tpl")
  sapply(sufs, lns)
  vfn <- paste(fn, machext, sep="")
  file.rename(vfn, fn)
  rm(pkg.root, sufs, lns, machext, vfn)

  addat <- list(n=nrow(dat), pgr=ncol(mmat_gr), ptr=ncol(mmat_tr))
  addat <- c(addat, as.list(dat[, c("sumh", "tau", "k")]))
  addat <- c(addat, list(Zgr=mmat_gr, Ztr=mmat_tr))  
  R2admb:::dat_write(fn, addat)
  datn <- paste(fn, "dat", sep=".")
  cat("# Yij\n", file=datn, append=TRUE)
  tmpf <- function(x){
    x <- ifelse(x == 0, 1e-8, x)
    cat(x, "\n",  file=datn, append=TRUE)
  }
  invisible(lapply(datYij, tmpf))
  rm(datn)
  init <- list(ctr=rep(0, addat$ptr), cgr=rep(0, addat$pgr))
  if (fn == "epiglm") {
    init <- list(ctr=rep(0, addat$ptr), cgr=rep(0, addat$pgr))
    init$ctr[1] <- -5
    init$cgr[1] <- -0.4
  } else if (fn =="epiglmcan") {
    init$ctr[1] <- 0.006
    init$cgr[1] <- 0.6    
  }
  R2admb:::pin_write(fn, init)
  R2admb::run_admb(fn)
  ret <- try(R2admb::read_admb(fn))
  if(inherits(ret, "try-error")){
    ret$solution_type <- "admb_error"
  } else {
    ret$solution_type <- "admb_success"
  }
  ret
}

get_grad <- function(dat, mmat_gr, cgr, mmat_tr, ctr){
  betavec <- exp(mmat_tr %*% ctr)
  popvec <- dat$maxk + exp(mmat_gr %*% cgr)
  ret <- list()
  foo <- -popvec * dat$tau + (dat$tau * dat$k - dat$sumh) + dat$k / betavec
  foo <- foo * betavec
  ret$betagrad <- t(foo) %*% mmat_tr
  foo <- -betavec * dat$tau + digamma(popvec + 1) - digamma(popvec - dat$k + 1)
  foo <- foo * exp(mmat_gr %*% cgr)
  ret$popgrad <- t(foo) %*% mmat_gr
  ret 
}

get_admb_fit2 <- function(dat, datYij, mmat_gr, mmat_tr,
                         workdir=tempfile("admb_dir_")){
  ## So simultaneous ADMB runs cannot clobber each other's files
  dir.create(workdir, recursive=TRUE)
  original_dir <- getwd()
  setwd(workdir)
  on.exit(setwd(original_dir))

  fn <- "epiglm2"

  stopifnot(Sys.info()["sysname"] == "Linux")
  machext <- paste("-",   Sys.info()["machine"], sep="")
  pkg.root <- system.file(package="EpiGLM")
  lns <- function(suf){
    targ <- paste(fn, suf, sep="")
    src <- file.path(pkg.root, "admb", targ)
    file.copy(src, targ)    
  }
  sufs <- c(machext, ".tpl")
  sapply(sufs, lns)
  vfn <- paste(fn, machext, sep="")
  file.rename(vfn, fn)
  rm(pkg.root, sufs, lns, machext, vfn)

  addat <- list(cgr_phase=-1L, ctr_phase=1L, n=nrow(dat),
                pgr=ncol(mmat_gr), ptr=ncol(mmat_tr))
  addat <- c(addat, as.list(dat[, c("sumh", "tau", "k", "maxk")]))
  addat <- c(addat, list(Zgr=mmat_gr, Ztr=mmat_tr))

  R2admb:::dat_write(fn, addat)
  datn <- paste(fn, "dat", sep=".")
  cat("# Yij\n", file=datn, append=TRUE)
  catYij <- function(x){
    x <- ifelse(x == 0, 1e-8, x)
    cat(x, "\n",  file=datn, append=TRUE)
  }
  invisible(lapply(datYij, catYij))

  init <- list(ctr=rep(-1, addat$ptr), cgr=rep(0, addat$pgr))
  init$cgr[1] <- -10
  R2admb:::pin_write(fn, init)
  R2admb::run_admb(fn)
  ret <- try(R2admb::read_admb(fn))

  if(inherits(ret, "try-error")){
    ret$solution_type <- "admb_error_try1"
  } else {
    grad <- get_grad(dat=dat, mmat_gr=mmat_gr, cgr=init$cgr,
                     mmat_tr=mmat_tr, ctr=ret$coeflist$ctr)
    if (all(grad$popgrad < 0)) {
      ret$solution_type <- "minimum_feasible_X0"
    } else {
      R2admb::clean_admb(fn, "all")
      addat$cgr_phase <- 1L
      R2admb:::dat_write(fn, addat)
      cat("# Yij\n", file=datn, append=TRUE)
      invisible(lapply(datYij, catYij))
      init$cgr[1] <- -1
      init <- list(ctr=ret$coeflist$ctr, cgr=init$cgr)
      R2admb:::pin_write(fn, init)
      R2admb::run_admb(fn)
      ret <- try(R2admb::read_admb(fn))
      if(inherits(ret, "try-error")){
        ret$solution_type <- "admb_error_try2"
      } else {
        ret$solution_type <- "admb_solution2"
      }
    }
  }  
  ret
}

ipest2df <- function(job, res, alpha=0.05){
  propFreeImp <- structure(res$propFreeImp, names="propFreeImp")
  propOnsetNA <- structure(res$propOnsetNA, names="propOnsetNA")
  res <- res$ip_fit  
  hats <- coef(res$m)
  names(hats) <- paste("hat", names(hats), sep="_")
  ses <- try(res$coefs[, "Std. Error"])
  if (!is.null(ses)){
    z <- qnorm(1 - alpha / 2)
    lower <- structure(hats - z * ses, names="lower")
    upper <- structure(hats + z * ses, names="upper")
    width <- structure(upper - lower, names="width")
    disp <- structure(res$disp, names="ip_dispersion_hat")
    names(ses) <- paste("se", names(ses), sep="_")
    return(c(hats, propFreeImp, propOnsetNA, ses, lower, upper, width, disp))
  } else {
    return(c(hats, propFreeImp, propOnsetNA))
  }
}

trest2df <- function(job, res, alpha=0.05, method="admb", returnch=FALSE){
  fit <- res$fit[[method]]
  retch <- c(note=fit$solution_type, method=method)  
  ret <- c(propFreeImp=res$propFreeImp, propOnsetNA=res$propOnsetNA)
  if (!retch["note"] %in% c("admb_error", "no_transmission_data",
                            "divergent", "unidentifiable",
                            "exposure_rounded_to_zero") & returnch==FALSE){
    hats <- fit$coefficients
    names(hats) <- paste("hat", names(hats), sep="_")
    ret <- c(ret, hats)
    ## If hessian is positive definite
    if (retch["note"] %in% c("admb_success", "interior_point")){
      ses <- fit$se
      names(ses) <- paste("se", names(ses), sep="_")
      z <- qnorm(1 - alpha / 2)
      lower <- structure(hats - z * ses,
                         names=paste(names(hats), "lower", sep="_"))
      upper <- structure(hats + z * ses,
                         names=paste(names(hats), "upper", sep="_"))
      width <- structure(upper - lower,
                         names=paste(names(hats), "width", sep="_"))
      rho <- c(rho_ctr_gr=fit$cor["ctr", "cgr"])
      ll <- c(logLik=fit$loglik)
      ret <- c(ret, ses, lower, upper, width, rho, ll)
    }
  }
  if(returnch) retch else ret  
}

regest2df <- function(job, res, alpha=0.05, returnch=FALSE){
  fit <- res$fit$admb
  retch <- c(note=fit$solution_type)  
  ret <- c(propFreeImp=res$propFreeImp)
  if (!retch["note"] %in% c("admb_error", "no_transmission_data",
                            "divergent", "unidentifiable",
                            "exposure_rounded_to_zero") & returnch==FALSE){
    hats <- fit$coefficients
    names(hats) <- paste("hat", names(hats), sep="_")
    ret <- c(ret, hats)
    ## If hessian is positive definite
    if (retch["note"] %in% c("admb_success", "interior_point")){
      ses <- fit$se
      names(ses) <- paste("se", names(ses), sep="_")
      z <- qnorm(1 - alpha / 2)
      lower <- structure(hats - z * ses,
                         names=paste(names(hats), "lower", sep="_"))
      upper <- structure(hats + z * ses,
                         names=paste(names(hats), "upper", sep="_"))
      width <- structure(upper - lower,
                         names=paste(names(hats), "width", sep="_"))
      ll <- c(logLik=fit$loglik)
      ret <- c(ret, ses, lower, upper, width, ll)
    }
  }
  if(returnch) retch else ret  
}

est2df <- function(job, res, alpha=0.05, returnch=FALSE){
  fit <- res$fit$admb
  ipfit <- res$ip_fit
  retch <- c(note=fit$solution_type)  
  ret <- c(propFreeImp=res$propFreeImp)
  if (!retch["note"] %in% c("admb_error", "no_transmission_data",
                            "divergent", "unidentifiable",
                            "exposure_rounded_to_zero") & returnch==FALSE){
    hats <- fit$coefficients
    if(!inherits(ipfit$m, "try-error")){
      hats <- c(hats, coef(ipfit$m))
      ipses <- try(ipfit$coefs[, "Std. Error"])      
    }
    names(hats) <- paste("hat", names(hats), sep="_")
    ret <- c(ret, hats)
    ## If hessian is positive definite
    if (retch["note"] %in% c("admb_success", "interior_point")){
      ses <- fit$se
      if(!inherits(ipfit, "try-error")){
        ses <- c(ses, ipfit$coefs[, "Std. Error"])
      }      
      names(ses) <- paste("se", names(ses), sep="_")
      z <- qnorm(1 - alpha / 2)
      lower <- structure(hats - z * ses,
                         names=paste(names(hats), "lower", sep="_"))
      upper <- structure(hats + z * ses,
                         names=paste(names(hats), "upper", sep="_"))
      width <- structure(upper - lower,
                         names=paste(names(hats), "width", sep="_"))
      ll <- c(logLik=fit$loglik)
      ret <- c(ret, ses, lower, upper, width, ll)
    }
  }
  if(returnch) retch else ret  
}
