///|
/// Stage-efficiency allocation and capacity planning utilities.
pub(all) struct StageAllocation {
  stage : Int
  ideal_efficiency : Double
  allocated_efficiency : Double
  expected_transfer : Double
  load_fraction : Double
  adjustment : Double
} derive(Debug, Eq, ToJson)

///|
pub(all) struct AllocationPlan {
  stages : Array[StageAllocation]
  average_efficiency : Double
  minimum_efficiency : Double
  maximum_efficiency : Double
  total_transfer : Double
  balanced : Bool
} derive(Debug, Eq, ToJson)

///|
pub(all) struct EfficiencyScenario {
  name : String
  multiplier : Double
  recovery : Double
  required_stages : Int
} derive(Debug, Eq, ToJson)

///|
pub fn validate_efficiencies(efficiencies : ArrayView[Double]) -> Bool {
  if efficiencies.length() == 0 {
    return false
  }
  for efficiency in efficiencies {
    if efficiency <= 0.0 || efficiency > 1.0 {
      return false
    }
  }
  true
}

///|
pub fn allocate_by_load(
  ideal_efficiencies : ArrayView[Double],
  stage_loads : ArrayView[Double],
  total_efficiency : Double,
) -> AllocationPlan raise @domain.ExtractionError {
  if ideal_efficiencies.length() == 0 ||
    ideal_efficiencies.length() != stage_loads.length() ||
    !validate_efficiencies(ideal_efficiencies) ||
    total_efficiency <= 0.0 ||
    total_efficiency > ideal_efficiencies.length().to_double() {
    raise @domain.ExtractionError::InvalidArgument(
      message="efficiency allocation inputs are invalid",
    )
  }
  let mut load_total = 0.0
  for load in stage_loads {
    if load < 0.0 {
      raise @domain.ExtractionError::InvalidArgument(
        message="stage load cannot be negative",
      )
    }
    load_total += load
  }
  if load_total == 0.0 {
    load_total = ideal_efficiencies.length().to_double()
  }
  let allocations = []
  let mut sum_efficiency = 0.0
  let mut minimum = 1.0e30
  let mut maximum = 0.0
  let mut transfer = 0.0
  for index in 0.. AllocationPlan raise @domain.ExtractionError {
  if stage_count < 1 || efficiency <= 0.0 || efficiency > 1.0 {
    raise @domain.ExtractionError::InvalidArgument(
      message="uniform allocation inputs are invalid",
    )
  }
  let efficiencies = Array::make(stage_count, efficiency)
  let loads = Array::make(stage_count, 1.0)
  allocate_by_load(efficiencies, loads, efficiency * stage_count.to_double())
}

///|
pub fn plan_total_efficiency(plan : AllocationPlan) -> Double {
  let mut result = 1.0
  for stage in plan.stages {
    result *= stage.allocated_efficiency
  }
  result
}

///|
pub fn plan_stage(
  plan : AllocationPlan,
  stage : Int,
) -> StageAllocation raise @domain.ExtractionError {
  if stage < 1 || stage > plan.stages.length() {
    raise @domain.ExtractionError::OutOfRange(
      name="stage",
      value=stage.to_double(),
      lower=1.0,
      upper=plan.stages.length().to_double(),
    )
  }
  plan.stages[stage - 1]
}

///|
pub fn plan_rebalance(
  plan : AllocationPlan,
  target_efficiency : Double,
) -> AllocationPlan raise @domain.ExtractionError {
  if target_efficiency <= 0.0 || target_efficiency > 1.0 {
    raise @domain.ExtractionError::InvalidFraction(
      name="target efficiency",
      value=target_efficiency,
    )
  }
  let ideals = Array::make(plan.stages.length(), target_efficiency)
  let loads = [ for stage in plan.stages => stage.load_fraction ]
  allocate_by_load(
    ideals,
    loads,
    target_efficiency * plan.stages.length().to_double(),
  )
}

///|
pub fn allocation_adjustments(plan : AllocationPlan) -> Array[Double] {
  [
    for stage in plan.stages => stage.adjustment
  ]
}

///|
pub fn allocation_transfer(plan : AllocationPlan) -> Array[Double] {
  [
    for stage in plan.stages => stage.expected_transfer
  ]
}

///|
pub fn allocation_max_adjustment(plan : AllocationPlan) -> Double {
  let mut result = 0.0
  for stage in plan.stages {
    result = result.max(stage.adjustment.abs())
  }
  result
}

///|
pub fn allocation_is_stable(
  plan : AllocationPlan,
  maximum_adjustment : Double,
) -> Bool {
  plan.balanced && allocation_max_adjustment(plan) <= maximum_adjustment
}

///|
pub fn required_stages_for_efficiency(
  target_recovery : Double,
  stage_efficiency : Double,
) -> Int raise @domain.ExtractionError {
  if target_recovery <= 0.0 ||
    target_recovery >= 1.0 ||
    stage_efficiency <= 0.0 ||
    stage_efficiency > 1.0 {
    raise @domain.ExtractionError::InvalidArgument(
      message="stage requirement inputs are invalid",
    )
  }
  let mut stages = 1
  let mut recovery = stage_efficiency
  while recovery < target_recovery && stages < 100000 {
    stages += 1
    recovery = recovery_after_stages(stage_efficiency, stages)
  }
  stages
}

///|
pub fn scenario_requirements(
  names : ArrayView[String],
  multipliers : ArrayView[Double],
  target_recovery : Double,
  base_efficiency : Double,
) -> Array[EfficiencyScenario] raise @domain.ExtractionError {
  let count = names.length().min(multipliers.length())
  let scenarios = []
  for index in 0.. EfficiencyScenario raise @domain.ExtractionError {
  if scenarios.length() == 0 {
    raise @domain.ExtractionError::InsufficientData(required=1, actual=0)
  }
  let mut selected = scenarios[0]
  for scenario in scenarios {
    if scenario.recovery < selected.recovery ||
      (
        scenario.recovery == selected.recovery &&
        scenario.required_stages > selected.required_stages
      ) {
      selected = scenario
    }
  }
  selected
}

///|
pub fn allocation_csv(plan : AllocationPlan) -> String {
  let lines = [
    "stage,ideal_efficiency,allocated_efficiency,expected_transfer,load_fraction,adjustment",
  ]
  for stage in plan.stages {
    lines.push(
      stage.stage.to_string() +
      "," +
      stage.ideal_efficiency.to_string() +
      "," +
      stage.allocated_efficiency.to_string() +
      "," +
      stage.expected_transfer.to_string() +
      "," +
      stage.load_fraction.to_string() +
      "," +
      stage.adjustment.to_string(),
    )
  }
  lines.join("\n")
}

///|
pub fn allocation_markdown(plan : AllocationPlan) -> String {
  let lines = [
    "| Stage | Ideal efficiency | Allocated | Expected transfer | Load fraction | Adjustment |",
    "|---:|---:|---:|---:|---:|---:|",
  ]
  for stage in plan.stages {
    lines.push(
      "| " +
      stage.stage.to_string() +
      " | " +
      stage.ideal_efficiency.to_string() +
      " | " +
      stage.allocated_efficiency.to_string() +
      " | " +
      stage.expected_transfer.to_string() +
      " | " +
      stage.load_fraction.to_string() +
      " | " +
      stage.adjustment.to_string() +
      " |",
    )
  }
  lines.join("\n")
}

///|
pub fn scenario_markdown(scenarios : ArrayView[EfficiencyScenario]) -> String {
  let lines = [
    "| Scenario | Multiplier | Recovery | Required stages |", "|---|---:|---:|---:|",
  ]
  for scenario in scenarios {
    lines.push(
      "| " +
      scenario.name +
      " | " +
      scenario.multiplier.to_string() +
      " | " +
      scenario.recovery.to_string() +
      " | " +
      scenario.required_stages.to_string() +
      " |",
    )
  }
  lines.join("\n")
}

///|
fn recovery_after_stages(efficiency : Double, stages : Int) -> Double {
  let mut remaining = 1.0
  for _ in 0..