
#' @title Save dcm Data to Output List Format
#'
#' @description This function transfers the data from the main `df`
#'              object to the output `out` object at the end of each
#'              run in [dcm()].
#'
#' @param df Main object in [dcm()] simulations.
#' @param s Current run number.
#' @param param Param list set in [param.dcm()].
#' @param control Control list set in [control.dcm()].
#' @param out Out list passed back in for updating at runs 2+.
#'
#' @return A list with the following elements:
#'
#'  * **param:** the epidemic parameters passed into the model through
#'        [param.dcm()], with additional parameters added as
#'        necessary.
#'  * **control:** the control settings passed into the model through
#'        [control.dcm()], with additional controls added as
#'        necessary.
#'  * **epi:** a list of data frames, one for each epidemiological
#'        output from the model.
#'
#'
#' @keywords internal
#' @export
#'
saveout.dcm <- function(df, s, param, control, out = NULL) {

  if (s == 1) {
    out <- list()
    out$param <- param
    out$control <- control
    if (length(control$nsteps) == 1) {
      out$control$timesteps <- seq(1, control$nsteps, control$dt)
    } else {
      out$control$timesteps <- control$nsteps
    }

    out$epi <- list()
    for (j in 2:ncol(df)) {
      out$epi[[names(df)[j]]] <- data.frame(df[, j])
    }
  } else {
    for (j in 2:ncol(df)) {
      out$epi[[names(df)[j]]][, s] <- data.frame(df[, j])
    }
  }

  # Remove NA's from flows by setting last value to penultimate value
  ns <- nrow(out$epi[[1]])
  lr.na <- sapply(out$epi, function(x) is.na(x[ns, s]) & !is.na(x[ns - 1, s]))
  wh.lr.na <- as.numeric(which(lr.na == TRUE))
  if (length(wh.lr.na) > 0) {
    for (jj in wh.lr.na) {
      out$epi[[jj]][ns, s] <- out$epi[[jj]][ns - 1, s]
    }
  }


  ## Processing for final run
  if (s == control$nruns) {

    # Set column names for varying list elements
    for (s in as.vector(which(lapply(out$epi, class) == "data.frame"))) {
      colnames(out$epi[[s]]) <- paste0("run", 1:control$nruns)
    }

  }

  return(out)
}


#' @title Save icm Data to Output List Format
#'
#' @description This function transfers the data from the main `icm_dat`
#'              class data object to the output `out` object at the end of
#'              each simulation in [icm()].
#'
#' @inheritParams prevalence.icm
#' @param s Current simulation number.
#' @param out Out list passed back in for updating at simulations 2+.
#'
#' @return
#' A list with the following elements:
#'
#'  * **param:** the epidemic parameters passed into the model through
#'        [param.icm()], with additional parameters added as
#'        necessary.
#'  * **control:** the control settings passed into the model through
#'        [control.icm()], with additional controls added as
#'        necessary.
#'  * **epi:** a list of data frames, one for each epidemiological
#'        output from the model.
#'
#'
#' @keywords internal
#' @export
#'
saveout.icm <- function(dat, s, out = NULL) {

  if (s == 1) {
    out <- list()
    out$param <- dat$param
    out$control <- dat$control
    out$epi <- list()
    for (j in seq_along(dat$epi)) {
      out$epi[[names(dat$epi)[j]]] <- data.frame(dat$epi[j])
    }
  } else {
    for (j in seq_along(dat$epi)) {
      out$epi[[names(dat$epi)[j]]][, s] <- data.frame(dat$epi[j])
    }
  }

  ## Processing for final run
  if (s == dat$control$nsims) {

    # Remove functions from control list
    ftodel <- grep(".FUN", names(out$control), value = TRUE)
    out$control[ftodel] <- NULL

    # Set column names for varying list elements
    for (i in as.vector(which(lapply(out$epi, class) == "data.frame"))) {
      colnames(out$epi[[i]]) <- paste0("sim", 1:dat$control$nsims)
    }

  }

  return(out)
}


#' @title Save netsim Data to Output List Format
#'
#' @description This function transfers the data from the main `netsim_dat`
#'              object to the output `out` object at the end of each
#'              simulation in [netsim()].
#'
#' @inheritParams recovery.net
#' @param s Current simulation number.
#' @param out Out list passed back in for updating at simulations 2+.
#'
#' @return
#' A list with the following elements:
#'
#'  * **param:** the epidemic parameters passed into the model through
#'        [param.net()], with additional parameters added as
#'        necessary.
#'  * **control:** the control settings passed into the model through
#'        [control.net()], with additional controls added as
#'        necessary.
#'  * **epi:** a list of data frames, one for each epidemiological
#'        output from the model.
#'  * **stats:** a list containing two sublists, `nwstats` for any
#'        network statistics saved in the simulation, and `transmat` for
#'        the transmission matrix saved in the simulation.
#'  * **network:** a list of `networkDynamic` objects,
#'         one for each model simulation.
#'
#'
#' @keywords internal
#' @export
#'
saveout.net <- function(dat, s, out = NULL) {

  ## Shared elements: only set up once on first simulation
  if (s == 1) {
    out <- list()
    out$param <- dat$param
    out$control <- dat$control
    out$nwparam <- dat$nwparam
    out$num.nw <- dat$num.nw

    out$epi <- list()
    for (j in seq_along(dat$epi)) {
      out$epi[[names(dat$epi)[j]]] <- data.frame(dat$epi[j])
    }
  }

  ## Shared elements: merge random param values for s > 1
  if (s > 1) {
    if (!is.null(dat$param$random.params.values)) {
      for (nms in names(dat$param$random.params.values)) {
        if (length(dat$param$random.params.values[[nms]]) > 1) {
          if (!is.list(out$param$random.params.values[[nms]])) {
            out$param$random.params.values[[nms]] <- list(
              out$param$random.params.values[[nms]]
            )
          }

          out$param$random.params.values[[nms]] <- c(
            out$param$random.params.values[[nms]],
            list(dat$param$random.params.values[[nms]])
          )

        } else {
          out$param$random.params.values[[nms]] <- c(
            out$param$random.params.values[[nms]],
            dat$param$random.params.values[[nms]]
          )
        }
      }
    }

    for (j in seq_along(dat$epi)) {
      out$epi[[names(dat$epi)[j]]][, s] <- data.frame(dat$epi[j])
    }
  }

  ## Per-simulation elements: use [s] <- list() to preserve NULLs
  out$coef.form[s] <- list(lapply(dat$nwparam, `[[`, "coef.form"))

  if (dat$control$save.run) {
    out$run[s] <- list(dat$run)
  }

  if (dat$control$save.cumulative.edgelist) {
    out$cumulative.edgelist[s] <- list(get_cumulative_edgelists_df(dat))
  }

  out$attr.history[s] <- list(dat$attr.history)
  out$raw.records[s] <- list(dat$raw.records)

  if (dat$control$save.nwstats == TRUE) {
    ## bind rows
    nwstats <- lapply(dat$stats$nwstats, dplyr::bind_rows)
    ## compute ess
    nwstats <- lapply(nwstats, function(y) structure(y, ess = ess(y)))
    out$stats$nwstats[s] <- list(nwstats)
  }

  if (dat$control$save.transmat == TRUE) {
    if (!is.null(dat$stats$transmat)) {
      transmat <- dplyr::bind_rows(dat$stats$transmat)
      row.names(transmat) <- seq_len(nrow(transmat))
      out$stats$transmat[s] <- list(transmat)
    } else {
      out$stats$transmat[s] <- list(dplyr::tibble())
    }
    if (s == 1) {
      class(out$stats$transmat) <- c("transmat", class(out$stats$transmat))
    }
  }

  if (dat$control$save.network == TRUE) {
    ## call get_network to use most up-to-date el and attr in tergmLite case
    out$network[s] <- list(lapply(seq_len(dat$num.nw), get_network, x = dat))
  }

  if (!is.null(dat$control$save.other)) {
    for (i in seq_along(dat$control$save.other)) {
      el.name <- dat$control$save.other[i]
      if (el.name %in% names(dat)) {
        out[[el.name]][s] <- list(dat[[el.name]])
      } else if (el.name %in% names(dat$run)) {
        out[[el.name]][s] <- list(dat$run[[el.name]])
      } else {
        warning("`", el.name, "` is not saved in `dat` or `dat$run`")
      }
    }
  }

  if (dat$control$save.diss.stats == TRUE &&
        dat$control$save.network == TRUE &&
        dat$control$tergmLite == FALSE &&
        !is.null(dat$nwparam)) {

    ## for each simulated network, if dissolution model is edges-only, compute diss stats
    out$diss.stats[s] <- list(lapply(seq_len(dat$num.nw), function(network) {
      if (dat$nwparam[[network]]$coef.diss$diss.model.type == "edgesonly") {
        net <- get_network(dat, network = network)
        toggles_to_diss_stats(
          tedgelist_to_toggles(as.data.frame(net)),
          dat$nwparam[[network]]$coef.diss,
          dat$control$nsteps,
          net
        )
      } else {
        NULL
      }
    }))
  }

  ## Final processing
  if (s == dat$control$nsims) {

    # Set names for out
    simnames <- paste0("sim", seq_len(dat$control$nsims))
    for (i in as.vector(which(lapply(out$epi, class) == "data.frame"))) {
      colnames(out$epi[[i]]) <- simnames
    }

    if (length(out$run) > 0) {
      out$run <- name_saveout_elts(out$run, "run", simnames)
    }

    top_lvl_elts <- c("attr.history", ".records")
    for (elt in top_lvl_elts) {
      out[[elt]] <- name_saveout_elts(out[[elt]], elt, simnames)
    }

    if (dat$control$save.cumulative.edgelist) {
      out$cumulative.edgelist <- name_saveout_elts(
        out$cumulative.edgelist, "cumulative.edgelist", simnames
      )
    }

    out$coef.form <- name_saveout_elts(out$coef.form, "coef.form", simnames)

    if (dat$control$save.nwstats == TRUE) {
      out$stats$nwstats <- name_saveout_elts(out$stats$nwstats,
                                             "stats$nwstats", simnames)
    }

    if (dat$control$save.transmat == TRUE) {
      out$stats$transmat <- name_saveout_elts(out$stats$transmat,
                                              "stats$transmat", simnames)
    }

    if (dat$control$save.network == TRUE) {
      out$network <- name_saveout_elts(out$network, "network", simnames)
    }

    if (dat$control$save.diss.stats == TRUE &&
          dat$control$save.network == TRUE &&
          dat$control$tergmLite == FALSE) {
      out$diss.stats <- name_saveout_elts(out$diss.stats, "diss.stats", simnames)
    }

    # only rename top level elements, not the ones part of `run`
    for (el.name in dat$control$save.other) {
      if (el.name %in% names(out)) {
        out[[el.name]] <- name_saveout_elts(out[[el.name]], el.name, simnames)
      }
    }

    # Remove functions from control list
    ftodel <- grep(".FUN", names(out$control), value = TRUE)
    out$control[ftodel] <- NULL
    out$control$currsim <- NULL
    environment(out$control$nwstats.formula) <- NULL
  }

  return(out)
}


#' @title Save a List of netsim Data to Output List Format
#'
#' @description This function transfers the data from a list of the main
#'              `netsim_dat` objects to the output `out` object at the
#'              end of all simulations in [netsim()].
#'
#' @param dat_list A list of main `netsim_dat` objects in `netsim`
#'        simulations.
#'
#' @return
#' A list of class `netsim` with the following elements:
#'
#'  * **param:** the epidemic parameters passed into the model through
#'        `param`, with additional parameters added as necessary.
#'  * **control:** the control settings passed into the model through
#'        `control`, with additional controls added as necessary.
#'  * **epi:** a list of data frames, one for each epidemiological
#'        output from the model. Outputs for base models always include the
#'        size of each compartment, as well as flows in, out of, and between
#'        compartments.
#'  * **stats:** a list containing two sublists, `nwstats` for any
#'        network statistics saved in the simulation, and `transmat` for
#'        the transmission matrix saved in the simulation. See
#'        [control.net()] for further
#'        details.
#'  * **network:** a list of `networkDynamic` objects,
#'         one for each model simulation.
#'
#'
#' @keywords internal
#' @export
#'
process_out.net <- function(dat_list) {
  for (s in seq_along(dat_list)) {
    # Set output
    if (s == 1) {
      out <- saveout.net(dat_list[[s]], s)
    } else {
      out <- saveout.net(dat_list[[s]], s, out)
    }
  }
  class(out) <- "netsim"

  return(out)
}

# Assign names to the save elements
# Gives a warning if there aren't enough elements
name_saveout_elts <- function(elt, elt_name, simnames) {
  if (length(elt) == length(simnames)) {
    names(elt) <- simnames
  } else if (length(elt) > 0) {
    warning("The number of `", elt_name, "` is not the number of simulations.",
            "Unable to assign each to the correct simulation.")
  }
  return(elt)
}
