
#' Predict smolt production
#'
#' @description Internal function that predicts the natural origin or hatchery origin smolt production from the
#' escapement, and saves state variables to [salmonMSE::salmonMSE_env].
#' - `smolt_func()` is the population dynamics function
#' - `makeRel_smolt()` generates a list for openMSE to use in the simulations
#'
#' @param Nage_NOS Matrix `[n_age, n_g]` of natural escapement from openMSE. `n_g` is the number of life history groups (sub groups within cohorts)
#' that contribute to spawning.
#' @param Nage_HOS Matrix `[n_age, n_r]` of hatchery escapement from openMSE. `n_r` is the number of release strategies (sub groups within cohorts)
#' that contribute to spawning.
#' @param Nage_stray Matrix `[nage, n_stray]` of hatchery escapement from donor populations that will stray to the recipient population. `stray_args$stray_matrix`
#' determines the proportions of the donor escapement that will move to the recipient population.
#' @param x Integer, simulation number from openMSE
#' @param y Integer, simulation year (including historical years)
#' @param output Character, whether to predict the natural origin or hatchery origin smolt production
#' @param s_enroute Numeric, en route survival of the escapement to the spawning grounds
#' @param p_female Vector `[maxage]`, proportion female for calculating the egg production.
#' @param fec Array `[nsim, maxage, nyears+proyears]`. The fecundity at age schedule of spawners
#' @param SRRpars Data frame containing stock recruit parameters for natural smolt production from egg production.
#' Column names include: SRrel, kappa, capacity, Smax, phi
#' @param hatchery_args Named list containing various arguments controlling broodtake and hatchery production. See details below.
#' @param fitness_args Named list containing various arguments controlling population fitness from hatchery production.
#' Names include: fitness_type, omega2, theta, fitness_variance, heritability, zbar_start, fitness_floor, rel_loss
#' @param habitat_args Named list with arguments sontrolling freshwater life stage survival.
#' Names include: use_habitat (logical), Habitat (\linkS4class{Habitat} object)
#' @param stray_args Named list with arguments controlling strays in and out of the local population
#' Names include:
#' @param s Integer, salmonMSE population index. Used to report variables to [salmonMSE::salmonMSE_env].
#' @section hatchery_args:
#' Hatchery control parameters are included in a named list with the following arguments:
#'
#' - `egg_target` Numeric, target egg production for hatchery. Set to zero for no hatchery production.
#' - `ptarget_NOB` Numeric, the target proportion of the natural origin broodtake relative to the overall broodtake
#' - `pmax_NOB` Numeric, the maximum proportion of the natural origin escapement to be used as broodtake
#' - `fec_brood` Array `[nsim, maxage, nyears+proyears]`. the fecundity at age schedule of broodtake to calculate the total hatchery egg production
#' - `s_yearling` Numeric, the survival of eggs to the smolt life stage (for yearling release)
#' - `s_subyearling` Numeric. the survival of eggs to subyearling life stage (for subyearling release)
#' - `p_yearling` Numeric, the proportion of annual releases at the yearling life stage (vs. subyearling)
#' - `phatchery` Numeric, the proportion of the hatchery origin escapement that return to the hatchery,
#' for example, by removal from spawning grounds or swim-in facilities. These fish are available for broodtake.
#' - `premove_HOS` Numeric or function, the proportion of the hatchery origin fish to be removed from the spawning
#' grounds (in order to ensure a high proportion of NOS or for an in-river fishery). These fish are not available for broodtake.
#' For example, a value less than one can represent imperfect implementation of weir removal.
#' - `premove_NOS` Numeric or function, the proportion of the natural origin fish to be removed from the spawning
#' grounds, for example, an in-river fishery.
#' - `s_prespawn` Numeric, the survival of broodtake prior to egg production. `1 - s_prespawn` is the proportion of fish
#' not used for hatchery purposes, e.g., mortality or other resesarch purposes. Used to back-calculate the broodtake.
#' - `gamma` Numeric, the relative reproductive success of hatchery origin spawners (relative to natural origin spawners)
#' - `m` Numeric, mark rate for selective broodtake
#' @returns
#' - `smolt_func()` returns a numeric for the ratio of the realized smolt production vs. the hypothetical value if there were no
#' hatchery, en route mortality, or habitat improvement
#' - `makeRel_smolt()` returns a list that is passed to openMSE as a inter-population relationship
#' @keywords internal
smolt_func <- function(Nage_NOS, Nage_HOS, Nage_stray, x = -1, y, output = c("natural", "hatchery"),
                       s_enroute, p_female, fec, SRRpars, # Spawning (natural production)
                       hatchery_args, fitness_args, habitat_args, stray_args, s, g, prop_LHG, r) {

  if (y %% 2 != 0 || x < 0) return(0) # Odd time steps
  nyears <- 4
  yy <- max(floor((nyears-4)/2), 1)
  fec_t <- fec[x, , yy]
  if (!is.null(hatchery_args$fec_brood)) fec_brood_t <- hatchery_args$fec_brood[x, , yy]

  output <- match.arg(output)
  Nage_NOS[is.na(Nage_NOS)] <- 0

  # Hatchery fish
  if (missing(Nage_HOS)) Nage_HOS <- matrix(0, nrow(Nage_NOS), 1)
  Nage_HOS[is.na(Nage_HOS)] <- 0

  # Add incoming strays (external and movement from donor populations)
  stray_external <- stray_args$stray_external
  if (!missing(Nage_stray) && sum(Nage_stray)) {
    s_donor <- setdiff(stray_args$p_donor$s, s)
    for (ss in s_donor) {
      if (ss != s) {
        p <- stray_args$p_donor$p[stray_args$p_donor$s == ss]
        stray_external <- stray_external + stray_args$stray_matrix[ss, s] * Nage_stray[, paste0("Nage_", p)]
      }
    }
    # Remove donor strays
    Nage_HOS <- stray_args$stray_matrix[s, s] * Nage_HOS
  }

  # En-route mortality
  Nage_NOS_enroute <- Nage_NOS * s_enroute
  Nage_HOS_enroute <- Nage_HOS * s_enroute
  stray_external_enroute <- stray_external * s_enroute

  # Broodtake
  custom_brood_rule <- is.function(hatchery_args$f_brood) && !is.null(formals(hatchery_args$f_brood))
  if (hatchery_args$egg_target > 0 && sum(Nage_NOS, Nage_HOS)) {

    if (!custom_brood_rule) { # AHA approach
      Nage_NOS_avail_brood <- Nage_NOS_enroute * hatchery_args$pmax_esc
      Nage_HOS_avail_brood <- Nage_HOS_enroute * hatchery_args$pmax_esc
      broodtake <- calc_broodtake(
        NOR_escapement = Nage_NOS_avail_brood,
        HOR_escapement = Nage_HOS_avail_brood,
        stray_external_enroute,
        hatchery_args$brood_import,
        hatchery_args$ptarget_NOB,
        hatchery_args$pmax_NOB,
        hatchery_args$phatchery,
        hatchery_args$egg_target,
        hatchery_args$p_female,
        fec_brood_t,
        hatchery_args$s_prespawn,
        hatchery_args$m
      )
    } else { # Custom brood approach
      broodtake <- calc_broodtake_custom(
        f_brood = hatchery_args$f_brood,
        NOR_escapement = Nage_NOS_enroute,
        HOR_escapement = Nage_HOS_enroute,
        stray_external_enroute,
        hatchery_args$p_female,
        fec_brood_t,
        hatchery_args$s_prespawn,
        hatchery_args$m,
        hatchery_args$egg_target
      )
    }

    egg_NOB <- broodtake$egg_NOB
    egg_HOB <- broodtake$egg_HOB_unmarked + broodtake$egg_HOB_marked
    pNOB <- broodtake$pNOB

    hatchery_production <- calc_yearling(
      egg_NOB + egg_HOB,
      hatchery_args$s_yearling,
      hatchery_args$s_subyearling,
      hatchery_args$p_yearling,
      hatchery_args$p_subyearling
    )

    # List of preliminary spawners and in-river removals after broodtake
    spawners <- calc_spawners(
      broodtake, Nage_NOS_enroute, Nage_HOS_enroute, stray_external_enroute,
      ifelse(custom_brood_rule, NA, hatchery_args$phatchery),
      hatchery_args$premove_HOS, hatchery_args$premove_NOS, hatchery_args$m
    )

  } else {

    # Check for same format as list returned by .broodtake_func
    broodtake <- list(
      egg_NOB = 0,
      egg_HOB_unmarked = 0,
      egg_HOB_marked = 0,
      ptake_unmarked = 0,
      ptake_marked = 0,
      pNOB = 0,
      NOB = array(0, dim(Nage_NOS)),
      HOB_unmarked = array(0, dim(Nage_HOS)),
      HOB_marked = array(0, dim(Nage_HOS)),
      HOB_import = rep(0, nrow(Nage_NOS)),
      HOB_stray = array(0, dim(Nage_HOS))
    )
    egg_NOB <- egg_HOB <- pNOB <- 0

    hatchery_production <- list(
      yearling = rep(0, ncol(Nage_HOS)),
      subyearling = rep(0, ncol(Nage_HOS))
    )

    # Preliminary spawners after broodtake
    spawners <- calc_spawners(
      broodtake, Nage_NOS_enroute, Nage_HOS_enroute, stray_external_enroute,
      NA, # phatchery
      hatchery_args$premove_HOS, hatchery_args$premove_NOS, hatchery_args$m
    )

  }

  # Add habitat density-dependence on pre-spawn mortality
  if (habitat_args$use_habitat) {

    Habitat <- habitat_args$Habitat

    # By parental life cycle group x age
    NOS <- calc_SRR(
      spawners$NOS, sum(spawners$NOS, spawners$HOS),
      p = Habitat@prespawn_prod, capacity = Habitat@prespawn_capacity,
      type = Habitat@prespawn_rel
    )

    # By parental release strategy x age
    HOS <- calc_SRR(
      spawners$HOS, sum(spawners$NOS, spawners$HOS),
      p = Habitat@prespawn_prod, capacity = Habitat@prespawn_capacity,
      type = Habitat@prespawn_rel
    )

  } else {
    NOS <- spawners$NOS
    HOS <- spawners$HOS # Includes local and strays!
  }
  HOS_effective <- HOS * hatchery_args$gamma
  if (!sum(HOS) && is.null(hatchery_args$gamma)) HOS_effective <- HOS

  # Spawners weighted by fecundity
  pHOSeff <- sum(fec_t * HOS_effective)/sum(fec_t * NOS, fec_t * HOS_effective)
  pHOScensus <- sum(fec_t * HOS)/sum(fec_t * NOS, fec_t * HOS)

  # Natural egg production in the absence of fitness effects by life cycle group (NOS) and release strategy (HOS)
  Egg_NOS <- colSums(NOS * p_female * fec_t)
  Egg_HOS <- colSums(HOS_effective * p_female * fec_t)

  total_egg <- sum(Egg_NOS, Egg_HOS)

  # Fitness calculations
  if (x > 0 && nrow(salmonMSE_env$Ford) && any(fitness_args$fitness_type == "Ford")) {

    # Get zbar from salmonMSE_env
    zbar_prev <- filter(salmonMSE_env$Ford, .data$x == .env$x, .data$s == .env$s)

    if (nrow(zbar_prev)) {
      maxage <- nrow(NOS)
      zbar_brood <- matrix(0, maxage, 2) # Column 1 = natural environment, 2 = hatchery environment

      # Trait value by brood year
      for (a in 1:maxage) {
        if (sum(NOS[a, ]) || sum(broodtake$NOB[a, ])) {
          zbar1 <- dplyr::filter(zbar_prev, .data$t == .env$y - 2 * .env$a, .data$type == "natural") %>%
            pull(.data$zbar)
          if (!length(zbar1)) stop("Cannot find zbar") # zbar1 <- fitness_args$zbar_start[1]
          zbar_brood[a, 1] <- zbar1
        }
        if (sum(HOS_effective[a, ]) || sum(broodtake$HOB[a, ]) || sum(broodtake$HOB_import[a])) {
          zbar2 <- dplyr::filter(zbar_prev, .data$t == .env$y - 2 * .env$a, .data$type == "hatchery") %>%
            pull(.data$zbar)
          if (!length(zbar2)) stop("Cannot find zbar") # zbar2 <- fitness_args$zbar_start[2]
          zbar_brood[a, 2] <- zbar2
        }
      }

      zbar <- calc_zbar(
        rowSums(NOS), rowSums(HOS_effective), rowSums(broodtake$NOB),
        rowSums(broodtake$HOB_marked + broodtake$HOB_unmarked) + broodtake$HOB_import,
        fec_t, fec_brood_t, zbar_brood,
        fitness_args$fitness_variance, fitness_args$theta, fitness_args$phenotype_variance, fitness_args$heritability[x]
      )
    } else {
      zbar <- fitness_args$zbar_start
    }

    # Fitness in the natural and hatchery environments
    fitness <- rep(1, 2)

    if (sum(Nage_NOS, Nage_HOS)) {
      for (i in 1:2) {
        if (fitness_args$fitness_type[i] == "Ford") {
          fitness[i] <- calc_fitness(
            zbar[i], fitness_args$theta[i], fitness_args$fitness_variance,
            fitness_args$phenotype_variance, fitness_args$fitness_floor
          )
        }
      }
    }

    fitness_loss <- outer(fitness, fitness_args$rel_loss, "^")
  } else {
    zbar <- rep(100, 2)
    fitness <- c(1, 1)
    fitness_loss <- matrix(1, 2, 3)
  }

  # Egg production from egg output, then fry production (stage-specific density-dependent relationship as habitat variables)
  if (habitat_args$use_habitat) {
    t <- max(1, 0.5 * (y - habitat_args$nyears_om))
    Habitat <- habitat_args$Habitat

    ### Incubation survival
    # By parental life cycle group
    Egg_prod_NOS <- calc_SRR(
      Egg_NOS, sum(Egg_HOS, Egg_NOS),
      p = Habitat@egg_prod, capacity = Habitat@egg_capacity,
      type = Habitat@egg_rel
    )
    # By parental release strategy
    Egg_prod_HOS <- calc_SRR(
      Egg_HOS, sum(Egg_HOS, Egg_NOS),
      p = Habitat@egg_prod, capacity = Habitat@egg_capacity,
      type = Habitat@egg_rel
    )

    # By parental life cycle group
    Fry_NOS <- calc_SRR(
      Egg_prod_NOS, sum(Egg_prod_NOS, Egg_prod_HOS),
      p = Habitat@fry_prod * fitness_loss[1, 1], capacity = Habitat@fry_capacity * fitness_loss[1, 1],
      type = Habitat@fry_rel
    ) * Habitat@fry_sdev[x, t]

    # By parental release strategy
    Fry_HOS <- calc_SRR(
      Egg_prod_HOS, sum(Egg_prod_NOS, Egg_prod_HOS),
      p = Habitat@fry_prod * fitness_loss[1, 1], capacity = Habitat@fry_capacity * fitness_loss[1, 1],
      type = Habitat@fry_rel
    ) * Habitat@fry_sdev[x, t]

  } else {
    # Egg and fry production from egg output (by parental life cycle group/release strategy)
    Egg_prod_NOS <- Egg_NOS
    Egg_prod_HOS <- Egg_HOS

    Fry_NOS <- fitness_loss[1, 1] * Egg_prod_NOS
    Fry_HOS <- fitness_loss[1, 1] * Egg_prod_HOS
  }

  # Hatchery production after fitness loss by release strategy (next generation)
  yearling <- hatchery_production$yearling * fitness_loss[2, 1] * fitness_loss[2, 2]
  subyearling <- hatchery_production$subyearling * fitness_loss[2, 1]

  # Juveniles in freshwater environment together that experience density-dependent survival
  total_fry_DD <- sum(Fry_NOS, Fry_HOS)
  if (sum(yearling) && hatchery_args$yearling_DD) total_fry_DD <- sum(total_fry_DD, yearling)
  if (sum(subyearling) && hatchery_args$subyearling_DD) total_fry_DD <- sum(total_fry_DD, subyearling)

  # Smolt production
  SRrel <- match.arg(SRRpars[x, "SRrel"], choices = c("BH", "Ricker"))

  # Hatchery releases in competition with natural-origin juveniles with fitness effects (by release strategy)
  if (sum(yearling) && hatchery_args$yearling_DD) {
    if (habitat_args$use_habitat) {
      smolt_yearling <- calc_SRR(
        yearling, total_fry_DD,
        p = Habitat@smolt_prod * fitness_loss[2, 2], capacity = Habitat@smolt_capacity * fitness_loss[2, 2],
        type = Habitat@smolt_rel
      )
    } else {
      smolt_yearling <- calc_smolt(
        yearling, total_fry_DD,
        SRRpars[x, "kappa"], SRRpars[x, "capacity"], SRRpars[x, "Smax"], SRRpars[x, "phi"], SRRpars[x, "tau"],
        fitness_loss[2, 1] * fitness_loss[2, 2],
        SRrel
      )
    }
  } else {
    smolt_yearling <- yearling
  }

  if (sum(subyearling) && hatchery_args$subyearling_DD) {
    if (habitat_args$use_habitat) {
      smolt_subyearling <- calc_SRR(
        subyearling, total_fry_DD,
        p = Habitat@smolt_prod * fitness_loss[2, 2], capacity = Habitat@smolt_capacity * fitness_loss[2, 2],
        type = Habitat@smolt_rel
      )
    } else {
      smolt_subyearling <- calc_smolt(
        subyearling, total_fry_DD,
        SRRpars[x, "kappa"], SRRpars[x, "capacity"], SRRpars[x, "Smax"], SRRpars[x, "phi"], SRRpars[x, "tau"],
        fitness_loss[2, 1] * fitness_loss[2, 2],
        SRrel
      )
    }
  } else {
    smolt_subyearling <- subyearling
  }

  # Total outmigrating hatchery releases
  smolt_rel <- smolt_yearling + smolt_subyearling
  if (output == "hatchery") return(smolt_rel[r])

  # Natural fry production for life history group g (next generation)
  Fry_NOS_g <- sum(Fry_NOS) * prop_LHG
  Fry_HOS_g <- sum(Fry_HOS) * prop_LHG

  # Natural smolt production for life history group g (next generation)
  if (habitat_args$use_habitat) {
    smolt_NOS <- calc_SRR(
      Fry_NOS_g, total_fry_DD,
      p = Habitat@smolt_prod * fitness_loss[1, 2], capacity = Habitat@smolt_capacity * fitness_loss[1, 2],
      type = Habitat@smolt_rel
    ) * Habitat@smolt_sdev[x, t]
    smolt_HOS <- calc_SRR(
      Fry_HOS_g, total_fry_DD,
      p = Habitat@smolt_prod * fitness_loss[1, 2], capacity = Habitat@smolt_capacity * fitness_loss[1, 2],
      type = Habitat@smolt_rel
    ) * Habitat@smolt_sdev[x, t]
  } else {
    smolt_NOS <- calc_smolt(
      Fry_NOS_g, total_fry_DD,
      SRRpars[x, "kappa"], SRRpars[x, "capacity"], SRRpars[x, "Smax"], SRRpars[x, "phi"], SRRpars[x, "tau"],
      fitness_loss[1, 2],
      SRrel
    )
    smolt_HOS <- calc_smolt(
      Fry_HOS_g, total_fry_DD,
      SRRpars[x, "kappa"], SRRpars[x, "capacity"], SRRpars[x, "Smax"], SRRpars[x, "phi"], SRRpars[x, "tau"],
      fitness_loss[1, 2],
      SRrel
    )
  }

  # Predicted egg production naively calculated by openMSE
  # if there were no hatchery production, habitat improvement, enroute mortality, or multiple LHG
  Egg_openMSE <- sum(Nage_NOS[, g] * p_female * fec_t)

  # Predicted juveniles by openMSE
  if (habitat_args$use_habitat) {
    smolt_NOS_SRR <- 1
  } else {
    smolt_NOS_SRR <- calc_smolt(
      Egg_openMSE, Egg_openMSE,
      SRRpars[x, "kappa"], SRRpars[x, "capacity"], SRRpars[x, "Smax"], SRRpars[x, "phi"], SRRpars[x, "tau"],
      SRrel = SRrel
    )
  }

  # Perr_y is an adjustment to the predicted juveniles (as a result of everything happening in this function,
  # including brood, en-route mortality)
  #
  # MICE parameter (natural production, Perr = process error for recruitment of new cohort)
  # Updated here and passed to openMSE operating model so that brood year juveniles in marine life stage
  # in openMSE is equal to numbers predicted in this function (smolt_NOS + smolt_HOS)
  Perr_y <- (smolt_NOS[g] + smolt_HOS[g])/smolt_NOS_SRR
  if (!smolt_NOS_SRR) Perr_y[] <- 0

  # Write all state variables to salmonMSE_env
  if (output == "natural" && g == 1) {
    #no_hatchery <- hatchery_args$egg_target == 0
    #if (no_hatchery) r <- 1

    nr <- ncol(Nage_HOS)

    df_H <- lapply(1:nr, function(r) {
      data.frame(
        x = x,
        s = s,
        r = r,
        t = y, # Even time steps (remember MICE predicts Perr_y for next time step)
        a = 1:nrow(Nage_HOS),
        Esc_HOS = Nage_HOS[, r],
        HOB = broodtake$HOB_unmarked[, r] + broodtake$HOB_marked[, r] + broodtake$HOB_stray[, r],
        HOB_stray = broodtake$HOB_stray[, r],
        HOS = HOS[, r],
        HOS_stray = spawners$HOS_stray[, r],
        HOS_effective = HOS_effective[, r],
        HOB_import = if (r == 1) broodtake$HOB_import else rep(NA, nrow(Nage_HOS)),
        HOS_remove = spawners$HO_remove[, r]
      )
    }) %>%
      bind_rows()
    salmonMSE_env$H <- rbind(salmonMSE_env$H, df_H)

    df_stateH <- lapply(1:nr, function(r) {
      data.frame(
        x = x,
        s = s,
        r = r,
        t = y, # Even time steps (remember MICE predicts Perr_y for next time step)
        Egg_HOS = Egg_HOS[r], # Spawning output by RS r of this generation
        yearling = yearling[r],
        subyearling = subyearling[r],
        smolt_rel = smolt_rel[r]
      )
    }) %>%
      bind_rows()
    salmonMSE_env$stateH <- rbind(salmonMSE_env$stateH, df_stateH)

    ng <- ncol(Nage_NOS)
    df_N <- lapply(1:ng, function(g) {
      data.frame(
        x = x,
        s = s,
        g = g,
        t = y, # Even time steps (remember MICE predicts Perr_y for next time step)
        a = 1:nrow(Nage_NOS),
        Esc_NOS = Nage_NOS[, g],
        NOB = broodtake$NOB[, g],
        NOS = NOS[, g],
        NOS_remove = spawners$NO_remove[, g]
      )
    }) %>%
      bind_rows()
    salmonMSE_env$N <- rbind(salmonMSE_env$N, df_N)

    df_stateN <- lapply(1:ng, function(g) {
      data.frame(
        x = x,
        s = s,
        g = g,
        t = y, # Even time steps (remember MICE predicts Perr_y for next time step)
        Egg_NOS = Egg_NOS[g], # Egg production by LHG g of this generation
        Fry_NOS = Fry_NOS_g[g],  # Fry production assigned to LHG g for next generation
        Fry_HOS = Fry_HOS_g[g],
        smolt_NOS = smolt_NOS[g],
        smolt_HOS = smolt_HOS[g],
        pNOB = as.numeric(pNOB),
        pHOSeff = as.numeric(pHOSeff),
        pHOScensus = as.numeric(pHOScensus),
        Perr_y = Perr_y
      )
    }) %>%
      bind_rows()
    salmonMSE_env$stateN <- rbind(salmonMSE_env$stateN, df_stateN)

    # Save zbar
    df_Ford <- data.frame(
      x = x,
      s = s,
      t = y, # Even time steps (remember MICE predicts Perr_y for next time step)
      type = c("natural", "hatchery"),
      zbar = zbar,
      fitness = fitness
    )
    salmonMSE_env$Ford <- rbind(salmonMSE_env$Ford, df_Ford)
  }

  if (is.na(Perr_y)) stop("Perr_y is NA. Check calculations in smolt_func()")

  return(Perr_y)
}

#' @rdname smolt_func
#' @param p_smolt Integer, the population index for the smolt production in the openMSE model, corresponding to `output`
#' @param p_natural Integer vector, the population index for the natural origin escapement in the openMSE model. Can be more than one
#' if spawning is from multiple life history groups
#' @param p_hatchery Integer vector, the population index for the hatchery origin escapement in the openMSE model. Can be more than one
#' if there are multiple release strategies. Set to `NULL` for no hatchery production
#' @param p_stray Integer vector, population index for the hatchery strays in multi-system models
#' @param g Integer for the life history group of natural origin fish to pass the parameter back to openMSE (if `output = "natural"`)
#' @param prop_LHG Numeric vector, proportion of the egg production assign to life history groups (sums up to one) (only used if `output = "natural"`)
#' @param r Integer for the release strategy of hatchery origin fish to pass the parameter back to openMSE (if `output = "hatchery"`)
makeRel_smolt <- function(p_smolt = 1, s = 1, p_natural, p_hatchery = NULL,
                          p_stray = NULL,
                          maxage,
                          output = c("natural", "hatchery"), s_enroute,
                          p_female, fec, SRRpars,  # Spawning (natural production)
                          hatchery_args, fitness_args, habitat_args, stray_args, g, prop_LHG, r) {

  output <- match.arg(output)

  .smolt_func <- smolt_func

  formals(.smolt_func)$output <- output
  formals(.smolt_func)$s_enroute <- s_enroute
  formals(.smolt_func)$p_female <- p_female
  formals(.smolt_func)$fec <- fec

  formals(.smolt_func)$SRRpars <- SRRpars

  formals(.smolt_func)$hatchery_args <- hatchery_args
  formals(.smolt_func)$fitness_args <- fitness_args
  formals(.smolt_func)$habitat_args <- habitat_args
  formals(.smolt_func)$stray_args <- stray_args
  formals(.smolt_func)$s <- s

  if (output == "natural") {
    formals(.smolt_func)$g <- g
    formals(.smolt_func)$prop_LHG <- prop_LHG
  }
  if (output == "hatchery") formals(.smolt_func)$r <- r

  N_natural <- matrix(1, maxage, length(p_natural))

  if (!is.null(p_hatchery)) { # Need p_hatchery to use p_stray
    N_hatchery <- matrix(0, maxage, length(p_hatchery))
    N_stray <- matrix(0, maxage, length(p_stray))

    Perr_y <- .smolt_func(N_natural, N_hatchery, x = -1, y = 1)
    model <- c(Perr_y = Perr_y, colSums(N_natural), colSums(N_hatchery), colSums(N_stray), x = -1, y = 1)
    input <- paste0("Nage_", c(p_natural, p_hatchery, p_stray))

    natural_origin <- c(rep(TRUE, length(p_natural)), rep(FALSE, length(c(p_hatchery, p_stray))))
    stray <- c(rep(FALSE, length(c(p_natural, p_hatchery))), rep(TRUE, length(p_stray)))
  } else {
    Perr_y <- .smolt_func(N_natural, x = -1, y = 1)
    model <- c(Perr_y = Perr_y, colSums(N_natural), x = -1, y = 1)
    input <- paste0("Nage_", p_natural)

    natural_origin <- rep(TRUE, length(p_natural))
    stray <- rep(FALSE, length(p_natural))
  }

  response <- paste0("Perr_y_", p_smolt)
  terms <- c(response, input, "x", "y")

  out <- list(
    func = .smolt_func,
    model = structure(model, names = terms),
    fitted.values = Perr_y,
    CV = 0,
    terms = terms,
    lag = "next",
    type = "Recruitment deviation",
    Rel = switch(output,
                 "natural" = "Smolt natural production from escapement",
                 "hatchery" = "Smolt hatchery releases from escapement"),
    natural_origin = natural_origin,
    stray = stray
  )
  structure(out, class = "RelSmolt")
}

#' @importFrom stats predict
#' @export
predict.RelSmolt <- function(object, newdata, ...) {

  if (missing(newdata)) return(object$fitted.values)

  vars <- names(newdata)

  vars_Rel <- names(object$model)[-1]
  Nage <- vars_Rel[grepl("Nage", vars_Rel)]
  Nage_NOS <- Nage[object$natural_origin]
  Nage_HOS <- Nage[!object$natural_origin & !object$stray]
  if (any(object$stray)) Nage_stray <- Nage[object$stray]

  do_hatchery <- length(Nage_HOS) > 0

  val <- sapply(1:nrow(newdata), function(i) {
    Esc_NOS <- sapply(Nage_NOS, function(j) as.numeric(newdata[i, grepl(j, vars)])) # matrix n_age x ng
    n_age <- nrow(Esc_NOS)
    a <- seq(3, n_age, 2)
    if (sum(Esc_NOS[!a, ])) stop("Internal salmonMSE error: there is natural escapement in even age classes")

    if (do_hatchery) {
      Esc_HOS <- sapply(Nage_HOS, function(j) as.numeric(newdata[i, grepl(j, vars)]))
      if (sum(Esc_HOS[!a, ])) stop("Internal salmonMSE error: there is hatchery escapement in even age classes")
      if (any(object$stray)) {
        Esc_stray <- sapply(Nage_stray, function(j) as.numeric(newdata[i, grepl(j, vars)]))
        if (sum(Esc_stray[!a, ])) stop("Internal salmonMSE error: there are hatchery strays in even age classes")
      } else {
        Esc_stray <- matrix(0, n_age, 1)
      }
    } else {
      Esc_HOS <- Esc_stray <- matrix(0, n_age, 1)
    }
    x <- newdata[i, "x"]
    y <- newdata[i, "y"]
    object$func(Nage_NOS = Esc_NOS[a, , drop = FALSE], Nage_HOS = Esc_HOS[a, , drop = FALSE],
                Nage_stray = Esc_stray[a, , drop = FALSE], x = x, y = y)
  })

  return(val)

}

#' @importFrom stats simulate
#' @export
simulate.RelSmolt <- function(object, nsim = 1, seed = 1, ...) {

  set.seed(seed)
  stdev <- sqrt(log(1 + object$CV^2))
  do_sim <- object$fitted.values * rlnorm(nsim * length(object$fitted.values), -0.5 * stdev^2, stdev)
  val <- matrix(do_sim, length(object$fitted.values), nsim) %>% as.data.frame() %>%
    structure(names = paste0("sim_", 1:nsim), seed = seed)
  val
}

#' Update natural mortality of juveniles
#'
#' @description Internal function that updates juvenile natural mortality in the marine environement due
#' to fitness
#' - `SAR_fitness()` calculates the new natural mortality value
#' - `makeRel_SAR` generates a list for openMSE to use in the simulations
#' @param x Integer, simulation number from openMSE
#' @param y Integer, simulation year (including historical years)
#' @param envir Character, whether to obtain the fitness value for the natural or hatchery environment.
#' @param rel_loss Numeric, the loss exponent for the juveniles
#' @param s Integer, the salmonMSE population index. Used to search for the fitness value
#' @param nyears Integer, the number of historical years in the operating model
#' @param Mbase Array `[nsim, n_age, proyears]` the base natural mortality value in the openMSE operating model.
#' @returns
#' - `smolt_func()` returns a numeric for the ratio of the realized smolt production vs. the hypothetical value if there were no
#' hatchery, en route mortality, or habitat improvement
#' - `makeRel_smolt` returns a list that is passed to openMSE as a inter-population relationship
#' @keywords internal
SAR_fitness <- function(x = -1, y = 1,
                        envir = c("natural", "hatchery"),
                        rel_loss = 1, s = 1,
                        nyears, Mbase) {

  envir <- match.arg(envir)

  even_time_step <- !y %% 2
  if (even_time_step) {
    n_age <- dim(Mbase)[2]
    fitness_a <- rep(1, n_age)
    for (a in seq(2, n_age, 2)) {
      if (nrow(salmonMSE_env$Ford) && x > 0) {
        fitness_y <- dplyr::filter(
          salmonMSE_env$Ford,
          .data$x == .env$x, .data$s == .env$s, .data$t == y - a, .data$type == envir
        ) %>%
          pull(.data$fitness)

        if (length(fitness_y)) fitness_a[a] <- fitness_y
      }
    }

    fitness_loss <- fitness_a^rel_loss
    M <- Mbase[x, , y - nyears]
    SAR_base <- SAR_loss <- exp(-M)
    SAR_loss[M > .Machine$double.eps] <- SAR_base[M > .Machine$double.eps] * fitness_loss[M > .Machine$double.eps]
    M_loss <- -log(SAR_loss)
    if (any(is.na(M_loss))) stop("M_loss has NA. Check calculations in SAR_fitness()")
    return(M_loss)
  } else {
    return(Mbase[x, , y - nyears])
  }
}

#' @rdname SAR_fitness
#' @param p_smolt Integer, the population index for the juvenile population in the openMSE model
#' @keywords internal
makeRel_SAR <- function(p_smolt = 1, s = 1, envir = c("natural", "hatchery"),
                        rel_loss, nyears, Mbase) {

  envir <- match.arg(envir)

  .SAR_fitness <- SAR_fitness

  formals(.SAR_fitness)$s <- s
  formals(.SAR_fitness)$rel_loss <- rel_loss
  formals(.SAR_fitness)$envir <- envir
  formals(.SAR_fitness)$nyears <- nyears
  formals(.SAR_fitness)$Mbase <- Mbase

  model <- data.frame(M = 1, N = seq(0, 10), x = -1, y = 1)

  response <- paste0("M_", p_smolt)
  input <- paste0("N_", p_smolt)
  terms <- c(response, input, "x", "y")

  out <- list(
    func = .SAR_fitness,
    model = structure(model, names = terms),
    fitted.values = model$M,
    CV = 0,
    terms = terms,
    type = "Natural mortality",
    Rel = "Marine survival reduced by fitness",
    age = seq(0, dim(Mbase)[2] - 1)
  )
  structure(out, class = "SARfitness")
}

#' @export
predict.SARfitness <- function(object, newdata, ...) {
  if (missing(newdata)) return(object$fitted.values)
  val <- sapply(1:nrow(newdata), function(i) object$func(x = newdata[i, "x"], y = newdata[i, "y"]))
  return(val)
}

#' @export
simulate.SARfitness <- simulate.RelSmolt
