///|
pub(all) struct CrosscurrentStage {
  stage : Int
  solvent_flow : Double
  inlet_raffinate_concentration : Double
  outlet_raffinate_concentration : Double
  extract_concentration : Double
  extracted_solute : Double
  cumulative_recovery : Double
} derive(Debug, Eq, ToJson)

///|
pub(all) struct CrosscurrentResult {
  stages : Array[CrosscurrentStage]
  final_raffinate_concentration : Double
  total_extracted_solute : Double
  overall_recovery : Double
  balance_residual : Double
} derive(Debug, Eq, ToJson)

///|
/// Calculate ideal crosscurrent stages with fresh or preloaded solvent.
///
/// `solvent_flows` and `solvent_concentrations` contain one value per stage.
pub fn crosscurrent_constant_k(
  feed_flow~ : Double,
  feed_concentration~ : Double,
  solvent_flows : ArrayView[Double],
  solvent_concentrations? : Array[Double] = [],
  distribution~ : Double,
) -> CrosscurrentResult raise @domain.ExtractionError {
  @domain.require_positive("feed flow", feed_flow)
  @domain.require_non_negative("feed concentration", feed_concentration)
  @domain.require_non_negative("distribution coefficient", distribution)
  if solvent_flows.length() == 0 {
    raise @domain.ExtractionError::InsufficientData(required=1, actual=0)
  }
  if solvent_concentrations.length() != 0 &&
    solvent_concentrations.length() != solvent_flows.length() {
    raise @domain.ExtractionError::InvalidArgument(
      message="solvent concentrations must be empty or match solvent flows",
    )
  }
  let initial_solute = feed_flow * feed_concentration
  let stages : Array[CrosscurrentStage] = []
  let mut current = feed_concentration
  let mut total_extracted = 0.0
  let mut solvent_inlet_solute = 0.0
  for index, solvent_flow in solvent_flows {
    @domain.require_positive("solvent flow at stage \{index + 1}", solvent_flow)
    let solvent_concentration = if solvent_concentrations.length() == 0 {
      0.0
    } else {
      solvent_concentrations[index]
    }
    @domain.require_non_negative(
      "solvent concentration at stage \{index + 1}",
      solvent_concentration,
    )
    solvent_inlet_solute += solvent_flow * solvent_concentration
    let result = single_stage_constant_k(
      feed_flow~,
      feed_concentration=current,
      solvent_flow~,
      solvent_concentration~,
      distribution~,
    )
    total_extracted += result.extract_solute -
      solvent_flow * solvent_concentration
    let recovery = if initial_solute > 0.0 {
      total_extracted / initial_solute
    } else {
      0.0
    }
    stages.push({
      stage: index + 1,
      solvent_flow,
      inlet_raffinate_concentration: current,
      outlet_raffinate_concentration: result.raffinate_concentration,
      extract_concentration: result.extract_concentration,
      extracted_solute: result.extract_solute -
      solvent_flow * solvent_concentration,
      cumulative_recovery: recovery,
    })
    current = result.raffinate_concentration
  }
  let final_solute = feed_flow * current
  let residual = initial_solute +
    solvent_inlet_solute -
    final_solute -
    (total_extracted + solvent_inlet_solute)
  {
    stages,
    final_raffinate_concentration: current,
    total_extracted_solute: total_extracted,
    overall_recovery: if initial_solute > 0.0 {
      total_extracted / initial_solute
    } else {
      0.0
    },
    balance_residual: residual,
  }
}

///|
/// Equal-solvent allocation convenience wrapper.
pub fn crosscurrent_equal_solvent(
  feed_flow~ : Double,
  feed_concentration~ : Double,
  total_solvent~ : Double,
  stage_count~ : Int,
  distribution~ : Double,
) -> CrosscurrentResult raise @domain.ExtractionError {
  @domain.require_positive("total solvent", total_solvent)
  if stage_count <= 0 {
    raise @domain.ExtractionError::InvalidArgument(
      message="stage_count must be positive",
    )
  }
  let per_stage = total_solvent / stage_count.to_double()
  crosscurrent_constant_k(
    feed_flow~,
    feed_concentration~,
    Array::make(stage_count, per_stage),
    distribution~,
  )
}