///|
priv struct State {
  concentrations : Array[Double]
  activities : Array[Double]
  ionic_strength : Double
  charge_residual : Double
  alkalinity : Double
  element_residuals : Array[Double]
  residuals : Array[Double]
  max_scaled_residual : Double
  p_h : Double
}

///|
fn primary_species_indices(problem : Problem) -> Array[Int32] {
  let count = problem.elements.length() + 1
  let indices = Array::makei(count, fn(_) { -1 })
  for i = 0; i < problem.species.length(); i = i + 1 {
    match problem.species[i].base_component {
      Some(component) => indices[component] = i
      None => ()
    }
  }
  indices
}

///|
fn initial_guess(problem : Problem, alternate~ : Bool = false) -> Array[Double] {
  let primary = primary_species_indices(problem)
  let guess : Array[Double] = []
  for component = 0; component < primary.length(); component = component + 1 {
    let species_def = problem.species[primary[component]]
    let concentration = if component == problem.hydrogen_component {
      match problem.fixed_p_h {
        Some(p_h) => safe_exp(-p_h * private_log10)
        None => if alternate { 1.0e-3 } else { 1.0e-7 }
      }
    } else {
      let mut total = 1.0e-6
      for e = 0; e < problem.elements.length(); e = e + 1 {
        if species_def.element_stoich[e] > 0.0 {
          total = problem.elements[e].total_mol_l /
            species_def.element_stoich[e]
        }
      }
      total
    }
    let value = concentration.max(1.0e-30)
    let log_value = safe_log(value)
    guess.push(
      if alternate && component == problem.hydrogen_component {
        log_value + 3.0
      } else {
        log_value
      },
    )
  }
  guess
}

///|
fn evaluate_state(
  problem : Problem,
  gamma : Array[Double],
  variables : Array[Double],
) -> State {
  let primary = primary_species_indices(problem)
  let base_activities = Array::makei(primary.length(), fn(component) {
    let species_index = primary[component]
    let concentration = safe_exp(variables[component])
    concentration * gamma[species_index].max(1.0e-300)
  })
  let concentrations : Array[Double] = []
  let activities : Array[Double] = []
  for i = 0; i < problem.species.length(); i = i + 1 {
    let species_def = problem.species[i]
    let concentration = match species_def.base_component {
      Some(component) => safe_exp(variables[component])
      None => {
        let mut log_activity = species_def.log_k * private_log10
        for term in species_def.component_terms {
          log_activity = log_activity +
            term.coefficient * safe_log(base_activities[term.component])
        }
        safe_exp(log_activity) / gamma[i].max(1.0e-300)
      }
    }
    concentrations.push(concentration)
    activities.push(concentration * gamma[i].max(1.0e-300))
  }
  let element_totals = Array::makei(problem.elements.length(), fn(_) { 0.0 })
  let mut charge_residual = 0.0
  let mut alkalinity = 0.0
  for i = 0; i < problem.species.length(); i = i + 1 {
    let species_def = problem.species[i]
    for element = 0; element < problem.elements.length(); element = element + 1 {
      element_totals[element] = element_totals[element] +
        concentrations[i] * species_def.element_stoich[element]
    }
    charge_residual = charge_residual +
      concentrations[i] * species_def.charge.to_double()
    alkalinity = alkalinity + concentrations[i] * species_def.alkalinity
  }
  let element_residuals : Array[Double] = []
  let residuals : Array[Double] = []
  for element = 0; element < problem.elements.length(); element = element + 1 {
    let difference = element_totals[element] -
      problem.elements[element].total_mol_l
    element_residuals.push(difference)
    let scale = problem.elements[element].total_mol_l.abs().max(1.0e-300)
    residuals.push(difference / scale)
  }
  let control_residual = match (problem.fixed_p_h, problem.alkalinity_mol_l) {
    (Some(p_h), _) => {
      let hydrogen_index = primary[problem.hydrogen_component]
      safe_log(activities[hydrogen_index]) - -p_h * private_log10
    }
    (None, Some(target)) => {
      let scale = target.abs().max(1.0e-9)
      (alkalinity - target) / scale
    }
    (None, None) => {
      let mut charge_scale = 1.0e-9
      for i = 0; i < concentrations.length(); i = i + 1 {
        charge_scale = charge_scale +
          concentrations[i] * problem.species[i].charge.to_double().abs()
      }
      charge_residual / charge_scale.max(1.0e-9)
    }
  }
  residuals.push(control_residual)
  let mut maximum = 0.0
  for residual in residuals {
    maximum = maximum.max(residual.abs())
    if !finite(residual) {
      maximum = 1.0e300
    }
  }
  for value in variables {
    if !finite(value) || value.abs() > 690.0 {
      maximum = 1.0e300
    }
  }
  for value in concentrations {
    if !finite(value) || value <= 0.0 || value >= 1.0e290 || value <= 1.0e-290 {
      maximum = 1.0e300
    }
  }
  // Reject clipped mass-action exponentials, rather than report a clipped
  // concentration as a converged physical state.
  for i = 0; i < problem.species.length(); i = i + 1 {
    let species_def = problem.species[i]
    if species_def.base_component is None {
      let mut expected = species_def.log_k * private_log10
      for term in species_def.component_terms {
        expected = expected +
          term.coefficient * safe_log(base_activities[term.component])
      }
      if !finite(expected) || expected.abs() > 690.0 {
        maximum = 1.0e300
      }
    }
  }
  let ionic = ionic_strength_of(concentrations, problem.species)
  let hydrogen_index = primary[problem.hydrogen_component]
  let p_h = -safe_log10(activities[hydrogen_index])
  {
    concentrations,
    activities,
    ionic_strength: ionic,
    charge_residual,
    alkalinity,
    element_residuals,
    residuals,
    max_scaled_residual: maximum,
    p_h,
  }
}

///|
fn solve_linear_system(
  matrix : Array[Array[Double]],
  rhs : Array[Double],
) -> Array[Double]? {
  let n = rhs.length()
  if matrix.length() != n {
    return None
  }
  let a = Array::makei(n, fn(i) { matrix[i].copy() })
  let b = rhs.copy()
  for column = 0; column < n; column = column + 1 {
    let mut pivot = column
    let mut pivot_value = a[column][column].abs()
    for row = column + 1; row < n; row = row + 1 {
      let value = a[row][column].abs()
      if value > pivot_value {
        pivot = row
        pivot_value = value
      }
    }
    if pivot_value < 1.0e-14 {
      return None
    }
    if pivot != column {
      let row_copy = a[column]
      a[column] = a[pivot]
      a[pivot] = row_copy
      let b_copy = b[column]
      b[column] = b[pivot]
      b[pivot] = b_copy
    }
    let diagonal = a[column][column]
    for row = column + 1; row < n; row = row + 1 {
      let factor = a[row][column] / diagonal
      a[row][column] = 0.0
      for j = column + 1; j < n; j = j + 1 {
        a[row][j] = a[row][j] - factor * a[column][j]
      }
      b[row] = b[row] - factor * b[column]
    }
  }
  let solution = Array::makei(n, fn(_) { 0.0 })
  for row = n - 1; row >= 0; row = row - 1 {
    let mut value = b[row]
    for j = row + 1; j < n; j = j + 1 {
      value = value - a[row][j] * solution[j]
    }
    solution[row] = value / a[row][row]
  }
  Some(solution)
}

///|
fn inner_solve(
  problem : Problem,
  gamma : Array[Double],
  initial : Array[Double],
  options : SolveOptions,
) -> Result[(Array[Double], State, Int32), SolveError] {
  let mut variables = initial.copy()
  let mut state = evaluate_state(problem, gamma, variables)
  let n = variables.length()
  for iteration = 0
      iteration < options.max_inner_iterations
      iteration = iteration + 1 {
    if state.max_scaled_residual <= options.residual_tolerance {
      return Ok((variables, state, iteration))
    }
    let jacobian : Array[Array[Double]] = []
    for column = 0; column < n; column = column + 1 {
      let perturbed = variables.copy()
      perturbed[column] = perturbed[column] + options.finite_difference_step
      let shifted = evaluate_state(problem, gamma, perturbed)
      let column_values : Array[Double] = []
      for row = 0; row < n; row = row + 1 {
        column_values.push(
          (shifted.residuals[row] - state.residuals[row]) /
          options.finite_difference_step,
        )
      }
      jacobian.push(column_values)
    }
    // The matrix above is stored column-major while Gaussian elimination
    // consumes row-major values.
    let matrix = Array::makei(n, fn(row) {
      Array::makei(n, fn(column) { jacobian[column][row] })
    })
    let rhs = state.residuals.map(fn(value) { -value })
    let step = match solve_linear_system(matrix, rhs) {
      Some(value) => value
      None => return Err(Stalled("finite-difference Jacobian is singular"))
    }
    let mut accepted = false
    let mut damping = 1.0
    let current_norm = state.max_scaled_residual
    for trial = 0; trial < 14; trial = trial + 1 {
      let candidate = variables.copy()
      for i = 0; i < n; i = i + 1 {
        candidate[i] = candidate[i] + damping * step[i]
      }
      let candidate_state = evaluate_state(problem, gamma, candidate)
      if candidate_state.max_scaled_residual < current_norm {
        variables = candidate
        state = candidate_state
        accepted = true
        break
      }
      damping = damping * 0.5
    }
    if !accepted {
      return Err(
        Stalled("line search could not reduce the scaled balance residual"),
      )
    }
  }
  Err(Stalled("inner Newton iteration limit reached"))
}

///|
fn compare_solution(a : SpeciationResult, b : SpeciationResult) -> Bool {
  let delta_p_h = (a.p_h - b.p_h).abs()
  let mut delta_species = 0.0
  let count = a.species.length().min(b.species.length())
  for i = 0; i < count; i = i + 1 {
    let left = safe_log10(a.species[i].concentration_mol_l.max(1.0e-300))
    let right = safe_log10(b.species[i].concentration_mol_l.max(1.0e-300))
    delta_species = delta_species.max((left - right).abs())
  }
  delta_p_h > 1.0e-3 || delta_species > 1.0e-2
}

///|
fn saturation_results(
  problem : Problem,
  state : State,
) -> Array[SaturationResult] {
  let primary = primary_species_indices(problem)
  let base_activities = Array::makei(primary.length(), fn(component) {
    state.activities[primary[component]]
  })
  let results : Array[SaturationResult] = []
  for mineral in problem.minerals {
    let mut log_iap = 0.0
    for term in mineral.component_terms {
      log_iap = log_iap +
        term.coefficient * safe_log10(base_activities[term.component])
    }
    let si = log_iap - mineral.log_ksp
    let status = if si > 1.0e-6 {
      "OversaturatedCandidate"
    } else if si < -1.0e-6 {
      "Undersaturated"
    } else {
      "NearSaturation"
    }
    results.push({
      mineral: mineral.name,
      log_iap,
      log_ksp: mineral.log_ksp,
      saturation_index: si,
      state: status,
      source: mineral.source,
    })
  }
  results
}

///|
fn validate_options(options : SolveOptions) -> Result[Unit, SolveError] {
  if options.max_inner_iterations <= 0 || options.max_outer_iterations <= 0 {
    return Err(InvalidInput("iteration limits must be positive"))
  }
  if !finite(options.residual_tolerance) ||
    !finite(options.activity_tolerance) ||
    options.residual_tolerance <= 0.0 ||
    options.activity_tolerance <= 0.0 {
    return Err(InvalidInput("solver tolerances must be positive"))
  }
  if !finite(options.finite_difference_step) ||
    options.finite_difference_step <= 0.0 {
    return Err(InvalidInput("finite_difference_step must be positive"))
  }
  Ok(())
}

///|
fn solve_once(
  problem : Problem,
  options : SolveOptions,
  alternate~ : Bool = false,
) -> Result[SpeciationResult, SolveError] {
  match validate_options(options) {
    Ok(_) => ()
    Err(error) => return Err(error)
  }
  match validate_problem(problem) {
    Ok(_) => ()
    Err(error) => return Err(error)
  }
  let gamma = Array::makei(problem.species.length(), fn(_) { 1.0 })
  let mut variables = initial_guess(problem, alternate~)
  let mut total_inner = 0
  let trace : Array[IterationRecord] = []
  for outer = 0; outer < options.max_outer_iterations; outer = outer + 1 {
    let inner = match inner_solve(problem, gamma, variables, options) {
      Ok(value) => value
      Err(error) => return Err(error)
    }
    let (new_variables, state, inner_iterations) = inner
    variables = new_variables
    total_inner = total_inner + inner_iterations
    let next_gamma = match
      calculate_activity_coefficients(
        problem.activity_model,
        problem.temperature_k,
        state.ionic_strength,
        problem.species,
      ) {
      Ok(value) => value
      Err(error) => return Err(error)
    }
    let mut gamma_delta = 0.0
    for i = 0; i < gamma.length(); i = i + 1 {
      gamma_delta = gamma_delta.max(
        (safe_log10(next_gamma[i]) - safe_log10(gamma[i])).abs(),
      )
    }
    trace.push({
      outer_iteration: outer + 1,
      inner_iterations,
      ionic_strength: state.ionic_strength,
      max_scaled_residual: state.max_scaled_residual,
      max_log_gamma_error: gamma_delta,
    })
    if state.max_scaled_residual <= options.residual_tolerance &&
      gamma_delta <= options.activity_tolerance {
      let final_state = evaluate_state(problem, gamma, variables)
      if final_state.max_scaled_residual > options.residual_tolerance {
        continue
      }
      let diagnostics = [
        "Mass-action equations evaluated in log activity space".to_string(),
        "Element balances and the selected charge/pH/alkalinity control were checked".to_string(),
        "Mineral SI is diagnostic only; no precipitation phase switch was performed".to_string(),
        "Fixed alkalinity or fixed pH is a boundary condition; charge residual is retained as an audit value".to_string(),
        "25 C only; molal source constants applied with dilute molarity and unit water activity approximations".to_string(),
      ]
      if final_state.charge_residual.abs() > 1.0e-8 {
        diagnostics.push(
          "Nonzero charge residual: check counterions or account for the acid/base reservoir before interpreting as a complete water sample",
        )
      }
      let species_results : Array[SpeciesResult] = []
      for i = 0; i < problem.species.length(); i = i + 1 {
        species_results.push({
          name: problem.species[i].name,
          concentration_mol_l: final_state.concentrations[i],
          activity: final_state.activities[i],
          log_activity: safe_log10(final_state.activities[i]),
          activity_coefficient: gamma[i],
          charge: problem.species[i].charge,
          primary: problem.species[i].base_component is Some(_),
          source: problem.species[i].source,
        })
      }
      return Ok({
        status: "Converged",
        temperature_k: problem.temperature_k,
        volume_l: problem.volume_l,
        activity_model: problem.activity_model,
        p_h: final_state.p_h,
        ionic_strength: final_state.ionic_strength,
        charge_residual_mol_l: final_state.charge_residual,
        alkalinity_mol_l: final_state.alkalinity,
        element_residuals_mol_l: final_state.element_residuals,
        species: species_results,
        saturation: saturation_results(problem, final_state),
        activity_coefficients: gamma,
        inner_iterations: total_inner,
        outer_iterations: outer + 1,
        max_scaled_residual: final_state.max_scaled_residual,
        diagnostics,
        database_version: problem.database_version,
        solver_options: options,
        thermodynamics: problem.thermodynamics,
        iteration_trace: trace,
        max_log_gamma_error: gamma_delta,
      })
    }
    for i = 0; i < gamma.length(); i = i + 1 {
      gamma[i] = 0.5 * gamma[i] + 0.5 * next_gamma[i]
    }
  }
  Err(Stalled("outer activity self-consistency iteration limit reached"))
}

///|
/// Solve a validated finite aqueous system using damped finite-difference
/// Newton iterations inside a Davies or Debye-Hückel activity fixed point.
pub fn solve(
  problem : Problem,
  options~ : SolveOptions = default_options(),
) -> Result[SpeciationResult, SolveError] {
  let primary = solve_once(problem, options)
  let result = match primary {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  if options.check_multiple_solutions && problem.fixed_p_h is None {
    match solve_once(problem, options, alternate=true) {
      Ok(alternate) =>
        // A second converged branch is only suspicious if it differs by more
        // than the reporting precision.  This check is intentionally bounded.
        if compare_solution(result, alternate) {
          return Err(
            MultipleSolutionsSuspected(
              "alternate hydrogen initialisation reached a distinct branch",
            ),
          )
        }
      Err(error) =>
        result.diagnostics.push(
          "Alternate initialisation did not converge: " + error.message(),
        )
    }
  }
  Ok(result)
}