///|
pub(all) struct CountercurrentStage {
  stage : Int
  raffinate_concentration : Double
  extract_concentration : Double
  raffinate_solute_flow : Double
  extract_solute_flow : Double
} derive(Debug, Eq, ToJson)

///|
pub(all) struct CountercurrentResult {
  stages : Array[CountercurrentStage]
  raffinate_outlet_concentration : Double
  extract_outlet_concentration : Double
  recovery : Double
  balance_residual : Double
  extraction_factor : Double
} derive(Debug, Eq, ToJson)

///|
/// Solve an ideal countercurrent cascade with constant flows and constant K.
/// Stage one is adjacent to the feed; stage `N` is adjacent to fresh solvent.
pub fn countercurrent_constant_k(
  carrier_flow~ : Double,
  feed_concentration~ : Double,
  solvent_flow~ : Double,
  solvent_inlet_concentration? : Double = 0.0,
  distribution~ : Double,
  stage_count~ : Int,
) -> CountercurrentResult raise @domain.ExtractionError {
  @domain.require_positive("carrier flow", carrier_flow)
  @domain.require_positive("solvent flow", solvent_flow)
  @domain.require_non_negative("feed concentration", feed_concentration)
  @domain.require_non_negative(
    "solvent inlet concentration", solvent_inlet_concentration,
  )
  @domain.require_positive("distribution coefficient", distribution)
  if stage_count <= 0 {
    raise @domain.ExtractionError::InvalidArgument(
      message="stage_count must be positive",
    )
  }
  let vk = solvent_flow * distribution
  let diagonal = carrier_flow + vk
  let lower = Array::make(stage_count, -carrier_flow)
  let upper = Array::make(stage_count, -vk)
  let rhs = Array::make(stage_count, 0.0)
  lower[0] = 0.0
  upper[stage_count - 1] = 0.0
  rhs[0] = carrier_flow * feed_concentration
  rhs[stage_count - 1] += solvent_flow * solvent_inlet_concentration
  let c_prime = Array::make(stage_count, 0.0)
  let d_prime = Array::make(stage_count, 0.0)
  c_prime[0] = if stage_count > 1 { upper[0] / diagonal } else { 0.0 }
  d_prime[0] = rhs[0] / diagonal
  for index in 1.. {
      {
        stage: index + 1,
        raffinate_concentration: x,
        extract_concentration: distribution * x,
        raffinate_solute_flow: carrier_flow * x,
        extract_solute_flow: solvent_flow * distribution * x,
      }
    }
  ]
  let x_out = xs[stage_count - 1]
  let y_out = distribution * xs[0]
  let inlet = carrier_flow * feed_concentration +
    solvent_flow * solvent_inlet_concentration
  let outlet = carrier_flow * x_out + solvent_flow * y_out
  let removable_feed_solute = carrier_flow * feed_concentration
  {
    stages,
    raffinate_outlet_concentration: x_out,
    extract_outlet_concentration: y_out,
    recovery: if removable_feed_solute > 0.0 {
      carrier_flow * (feed_concentration - x_out) / removable_feed_solute
    } else {
      0.0
    },
    balance_residual: inlet - outlet,
    extraction_factor: vk / carrier_flow,
  }
}

///|
/// Find the smallest integer stage count that meets a raffinate target.
pub fn minimum_countercurrent_stages(
  carrier_flow~ : Double,
  feed_concentration~ : Double,
  solvent_flow~ : Double,
  solvent_inlet_concentration? : Double = 0.0,
  distribution~ : Double,
  target_raffinate~ : Double,
  maximum_stages? : Int = 200,
) -> CountercurrentResult raise @domain.ExtractionError {
  @domain.require_non_negative("target raffinate", target_raffinate)
  if target_raffinate >= feed_concentration {
    raise @domain.ExtractionError::InvalidArgument(
      message="target raffinate must be below feed concentration",
    )
  }
  if maximum_stages <= 0 {
    raise @domain.ExtractionError::InvalidArgument(
      message="maximum_stages must be positive",
    )
  }
  for stage_count in 1..<=maximum_stages {
    let result = countercurrent_constant_k(
      carrier_flow~,
      feed_concentration~,
      solvent_flow~,
      solvent_inlet_concentration~,
      distribution~,
      stage_count~,
    )
    if result.raffinate_outlet_concentration <= target_raffinate {
      return result
    }
  }
  raise @domain.ExtractionError::ConvergenceFailure(
    iterations=maximum_stages,
    residual=target_raffinate,
  )
}