///|
/// Stage result for a countercurrent cascade using a tabulated equilibrium curve.
pub(all) struct NonlinearCountercurrentStage {
  stage : Int
  raffinate_concentration : Double
  extract_concentration : Double
  raffinate_solute_flow : Double
  extract_solute_flow : Double
  local_residual : Double
} derive(Debug, Eq, ToJson)

///|
pub(all) struct NonlinearCountercurrentResult {
  stages : Array[NonlinearCountercurrentStage]
  raffinate_outlet_concentration : Double
  extract_outlet_concentration : Double
  recovery : Double
  balance_residual : Double
  iterations : Int
  converged : Bool
  maximum_local_residual : Double
} derive(Debug, Eq, ToJson)

///|
fn solve_stage_balance(
  carrier_flow : Double,
  solvent_flow : Double,
  incoming_solute : Double,
  equilibrium : @interpolation.EquilibriumTable,
  tolerance : Double,
  maximum_iterations : Int,
) -> Double raise @domain.ExtractionError {
  let lower = equilibrium.relation.lower()
  let upper = equilibrium.relation.upper()
  let residual = x => {
    carrier_flow * x + solvent_flow * equilibrium.y_at(x) - incoming_solute
  }
  let lower_residual = residual(lower)
  let upper_residual = residual(upper)
  if lower_residual > 0.0 || upper_residual < 0.0 {
    raise @domain.ExtractionError::OutOfRange(
      name="nonlinear stage balance",
      value=incoming_solute,
      lower=carrier_flow * lower + solvent_flow * equilibrium.y_at(lower),
      upper=carrier_flow * upper + solvent_flow * equilibrium.y_at(upper),
    )
  }
  let mut left = lower
  let mut right = upper
  let mut current = residual(left)
  for _ in 0.. tolerance * 10.0 {
    raise @domain.ExtractionError::ConvergenceFailure(
      iterations=maximum_iterations,
      residual=current,
    )
  }
  (left + right) / 2.0
}

///|
fn stage_residual(
  stage : Int,
  xs : ArrayView[Double],
  ys : ArrayView[Double],
  carrier_flow : Double,
  solvent_flow : Double,
  feed_concentration : Double,
  solvent_inlet_concentration : Double,
) -> Double {
  let upstream_x = if stage == 0 { feed_concentration } else { xs[stage - 1] }
  let downstream_y = if stage == xs.length() - 1 {
    solvent_inlet_concentration
  } else {
    ys[stage + 1]
  }
  carrier_flow * upstream_x +
  solvent_flow * downstream_y -
  carrier_flow * xs[stage] -
  solvent_flow * ys[stage]
}

///|
/// Solve a nonlinear countercurrent cascade by damped Gauss–Seidel iteration.
pub fn countercurrent_table(
  carrier_flow~ : Double,
  feed_concentration~ : Double,
  solvent_flow~ : Double,
  solvent_inlet_concentration? : Double = 0.0,
  equilibrium~ : @interpolation.EquilibriumTable,
  stage_count~ : Int,
  tolerance? : Double = 1.0e-9,
  maximum_iterations? : Int = 200,
  relaxation? : Double = 0.8,
) -> NonlinearCountercurrentResult 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("countercurrent tolerance", tolerance)
  @domain.require_positive("countercurrent relaxation", relaxation)
  if relaxation > 1.0 {
    raise @domain.ExtractionError::InvalidArgument(
      message="countercurrent relaxation must be at most one",
    )
  }
  if stage_count < 1 || maximum_iterations < 1 {
    raise @domain.ExtractionError::InvalidArgument(
      message="stage count and iteration limit must be positive",
    )
  }
  if feed_concentration < equilibrium.relation.lower() ||
    feed_concentration > equilibrium.relation.upper() {
    raise @domain.ExtractionError::OutOfRange(
      name="feed concentration",
      value=feed_concentration,
      lower=equilibrium.relation.lower(),
      upper=equilibrium.relation.upper(),
    )
  }
  let lower = equilibrium.relation.lower()
  let upper = equilibrium.relation.upper()
  let xs = Array::make(stage_count, 0.0)
  for index in 0.. tolerance * 100.0 {
    raise @domain.ExtractionError::ConvergenceFailure(
      iterations~,
      residual=maximum_residual,
    )
  }
  let stages = [
    for index, x in xs => {
      let residual = stage_residual(
        index, xs, ys, carrier_flow, solvent_flow, feed_concentration, solvent_inlet_concentration,
      )
      {
        stage: index + 1,
        raffinate_concentration: x,
        extract_concentration: ys[index],
        raffinate_solute_flow: carrier_flow * x,
        extract_solute_flow: solvent_flow * ys[index],
        local_residual: residual,
      }
    }
  ]
  let inlet = carrier_flow * feed_concentration +
    solvent_flow * solvent_inlet_concentration
  let outlet = carrier_flow * xs[stage_count - 1] + solvent_flow * ys[0]
  let feed_solute = carrier_flow * feed_concentration
  {
    stages,
    raffinate_outlet_concentration: xs[stage_count - 1],
    extract_outlet_concentration: ys[0],
    recovery: if feed_solute <= 0.0 {
      0.0
    } else {
      carrier_flow * (feed_concentration - xs[stage_count - 1]) / feed_solute
    },
    balance_residual: inlet - outlet,
    iterations,
    converged,
    maximum_local_residual: maximum_residual,
  }
}

///|
/// Find the first nonlinear cascade meeting a raffinate specification.
pub fn minimum_table_stages(
  carrier_flow~ : Double,
  feed_concentration~ : Double,
  solvent_flow~ : Double,
  solvent_inlet_concentration? : Double = 0.0,
  equilibrium~ : @interpolation.EquilibriumTable,
  target_raffinate~ : Double,
  maximum_stages? : Int = 100,
) -> NonlinearCountercurrentResult 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 < 1 {
    raise @domain.ExtractionError::InvalidArgument(
      message="maximum stage count must be positive",
    )
  }
  for count in 1..<=maximum_stages {
    let result = countercurrent_table(
      carrier_flow~,
      feed_concentration~,
      solvent_flow~,
      solvent_inlet_concentration~,
      equilibrium~,
      stage_count=count,
    )
    if result.raffinate_outlet_concentration <= target_raffinate {
      return result
    }
  }
  raise @domain.ExtractionError::ConvergenceFailure(
    iterations=maximum_stages,
    residual=target_raffinate,
  )
}

///|
/// Compare constant-K and tabulated equilibrium predictions.
pub(all) struct CascadeComparison {
  constant_result : CountercurrentResult
  table_result : NonlinearCountercurrentResult
  outlet_difference : Double
  recovery_difference : Double
} derive(Debug, Eq, ToJson)

///|
pub fn compare_countercurrent_models(
  carrier_flow~ : Double,
  feed_concentration~ : Double,
  solvent_flow~ : Double,
  distribution~ : Double,
  equilibrium~ : @interpolation.EquilibriumTable,
  stage_count~ : Int,
) -> CascadeComparison raise @domain.ExtractionError {
  let constant_result = countercurrent_constant_k(
    carrier_flow~,
    feed_concentration~,
    solvent_flow~,
    distribution~,
    stage_count~,
  )
  let table_result = countercurrent_table(
    carrier_flow~,
    feed_concentration~,
    solvent_flow~,
    equilibrium~,
    stage_count~,
  )
  {
    constant_result,
    table_result,
    outlet_difference: table_result.raffinate_outlet_concentration -
    constant_result.raffinate_outlet_concentration,
    recovery_difference: table_result.recovery - constant_result.recovery,
  }
}

///|
pub fn stage_residuals(result : NonlinearCountercurrentResult) -> Array[Double] {
  [
    for stage in result.stages => stage.local_residual
  ]
}

///|
pub fn maximum_concentration(result : NonlinearCountercurrentResult) -> Double {
  let mut maximum = 0.0
  for stage in result.stages {
    maximum = maximum.max(stage.raffinate_concentration)
    maximum = maximum.max(stage.extract_concentration)
  }
  maximum
}