///|
/// Termination state of a regularized discrete Wasserstein barycenter solve.
pub(all) enum BarycenterTermination {
  Converged
  IterationLimit
}

///|
/// The support weights, objective estimate, and convergence diagnostics of a
/// regularized barycenter.
pub(all) struct BarycenterResult {
  barycenter : Array[Double]
  weighted_objective : Double
  maximum_transport_residual : Double
  log_barycenter_residual : Double
  iterations : Int
  termination : BarycenterTermination
}

///|
/// Validation and numerical errors reported by `regularized_barycenter`.
pub(all) enum BarycenterError {
  EmptyMeasures
  WeightCountMismatch
  EmptySupport
  EmptyMeasure(Int)
  CostRowCountMismatch
  CostColumnCountMismatch(Int)
  NonFiniteWeight(Int)
  NonPositiveWeight(Int)
  NonFiniteWeightTotal
  NonFiniteMass(Int, Int)
  NegativeMass(Int, Int)
  NonFiniteMeasureTotal(Int)
  UnequalMeasureTotal(Int)
  NonFiniteCost(Int, Int)
  NegativeCost(Int, Int)
  InvalidEpsilon
  InvalidTolerance
  InvalidIterationLimit
  SolverError(SinkhornError)
  NumericalFailure
}

///|
/// Validation and numerical errors from `regularized_barycenter_with_costs`.
pub(all) enum BarycenterCostsError {
  InputError(BarycenterError)
  CostCountMismatch
  CostRowCountMismatch(Int)
  CostColumnCountMismatch(Int, Int)
  NonFiniteCost(Int, Int, Int)
  NegativeCost(Int, Int, Int)
  SolverError(Int, SinkhornError)
  NumericalFailure
}

///|
priv struct BarycenterValidatedInput {
  measures : Array[Array[Double]]
  scaled_weights : Array[Double]
  log_measures : Array[Array[Double]]
  support_size : Int
  total_mass : Double
}

///|
/// A shared matrix stays a single allocation for the legacy entry point;
/// heterogeneous matrices retain one matrix per input measure.
priv enum BarycenterCostMatrices {
  Shared(Array[Array[Double]])
  PerMeasure(Array[Array[Array[Double]]])
}

///|
priv enum BarycenterComputationError {
  SolverError(Int, SinkhornError)
  NumericalFailure
}

///|
/// Compute the fixed-support entropy-regularized Wasserstein barycenter.
///
/// Every input measure and the barycenter share one support with equal total
/// mass. Weights must be finite and strictly positive; they are normalized
/// internally. The common square cost matrix must be finite and nonnegative.
/// The iteration uses log-domain generalized Sinkhorn updates. Its tolerance
/// bounds the maximum change in normalized log barycenter between iterations.
/// The objective is the weighted sum of balanced Sinkhorn objectives, computed
/// after the barycenter iteration. The reported transport residual is the
/// largest marginal residual among those objective-evaluation solves.
pub fn regularized_barycenter(
  measures : Array[Array[Double]],
  weights : Array[Double],
  cost : Array[Array[Double]],
  options : SinkhornOptions,
) -> Result[BarycenterResult, BarycenterError] {
  let support_size = match
    barycenter_validate_header(measures, weights, options) {
    Err(error) => return Err(error)
    Ok(size) => size
  }
  match barycenter_validate_shared_cost(cost, support_size) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  let input = match
    barycenter_validate_measure_data(measures, weights, support_size) {
    Err(error) => return Err(error)
    Ok(validated) => validated
  }
  match
    barycenter_compute(input, BarycenterCostMatrices::Shared(cost), options) {
    Err(BarycenterComputationError::SolverError(_, error)) =>
      Err(BarycenterError::SolverError(error))
    Err(BarycenterComputationError::NumericalFailure) =>
      Err(BarycenterError::NumericalFailure)
    Ok(result) => Ok(result)
  }
}

///|
/// Compute a fixed-support regularized barycenter with one ground-cost matrix
/// per input measure.
///
/// `costs[k][i][j]` is the cost from common barycenter-support position `i` to
/// position `j` of `measures[k]`. Every matrix is square and has the same row
/// ordering, while its columns align with its paired measure. Costs need not be
/// symmetric or metric. The same epsilon and normalized positive weights apply
/// to all subproblems; this function does not learn or rescale costs.
pub fn regularized_barycenter_with_costs(
  measures : Array[Array[Double]],
  weights : Array[Double],
  costs : Array[Array[Array[Double]]],
  options : SinkhornOptions,
) -> Result[BarycenterResult, BarycenterCostsError] {
  let support_size = match
    barycenter_validate_header(measures, weights, options) {
    Err(error) => return Err(BarycenterCostsError::InputError(error))
    Ok(size) => size
  }
  match barycenter_validate_costs(costs, measures.length(), support_size) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  let input = match
    barycenter_validate_measure_data(measures, weights, support_size) {
    Err(error) => return Err(BarycenterCostsError::InputError(error))
    Ok(validated) => validated
  }
  match
    barycenter_compute(
      input,
      BarycenterCostMatrices::PerMeasure(costs),
      options,
    ) {
    Err(BarycenterComputationError::SolverError(index, error)) =>
      Err(BarycenterCostsError::SolverError(index, error))
    Err(BarycenterComputationError::NumericalFailure) =>
      Err(BarycenterCostsError::NumericalFailure)
    Ok(result) => Ok(result)
  }
}

///|
fn barycenter_validate_header(
  measures : Array[Array[Double]],
  weights : Array[Double],
  options : SinkhornOptions,
) -> Result[Int, BarycenterError] {
  let measure_count = measures.length()
  if measure_count == 0 {
    return Err(BarycenterError::EmptyMeasures)
  }
  if weights.length() != measure_count {
    return Err(BarycenterError::WeightCountMismatch)
  }
  if !is_finite(options.epsilon) || options.epsilon <= 0.0 {
    return Err(BarycenterError::InvalidEpsilon)
  }
  if !is_finite(options.tolerance) || options.tolerance <= 0.0 {
    return Err(BarycenterError::InvalidTolerance)
  }
  if options.max_iterations <= 0 {
    return Err(BarycenterError::InvalidIterationLimit)
  }
  let support_size = measures[0].length()
  if support_size == 0 {
    return Err(BarycenterError::EmptySupport)
  }
  Ok(support_size)
}

///|
fn barycenter_validate_shared_cost(
  cost : Array[Array[Double]],
  support_size : Int,
) -> Result[Unit, BarycenterError] {
  if cost.length() != support_size {
    return Err(BarycenterError::CostRowCountMismatch)
  }
  for i in 0.. Result[Unit, BarycenterCostsError] {
  if costs.length() != measure_count {
    return Err(BarycenterCostsError::CostCountMismatch)
  }
  for k in 0.. Result[BarycenterValidatedInput, BarycenterError] {
  let measure_count = measures.length()
  let mut maximum_weight = 0.0
  for k in 0.. maximum_weight {
      maximum_weight = weights[k]
    }
  }
  let scaled_weights = Array::make(measure_count, 0.0)
  let mut scaled_weight_total = 0.0
  for k in 0.. 0.0 {
        log_measure[i] = @math.ln(mass)
      }
    }
    if k == 0 {
      common_total = measure_total
    } else if measure_total > 0.0 || common_total > 0.0 {
      let scale = if measure_total > common_total {
        measure_total
      } else {
        common_total
      }
      if abs(measure_total - common_total) > 1.0e-12 * scale {
        return Err(BarycenterError::UnequalMeasureTotal(k))
      }
    }
    log_measures.push(log_measure)
  }
  Ok({
    measures,
    scaled_weights,
    log_measures,
    support_size,
    total_mass: common_total,
  })
}

///|
fn barycenter_zero_result(support_size : Int) -> BarycenterResult {
  {
    barycenter: Array::make(support_size, 0.0),
    weighted_objective: 0.0,
    maximum_transport_residual: 0.0,
    log_barycenter_residual: 0.0,
    iterations: 0,
    termination: BarycenterTermination::Converged,
  }
}

///|
fn barycenter_compute(
  input : BarycenterValidatedInput,
  costs : BarycenterCostMatrices,
  options : SinkhornOptions,
) -> Result[BarycenterResult, BarycenterComputationError] {
  if input.total_mass == 0.0 {
    return Ok(barycenter_zero_result(input.support_size))
  }
  let log_kernels = match barycenter_build_log_kernels(costs, options.epsilon) {
    Err(error) => return Err(error)
    Ok(kernels) => kernels
  }
  let log_scalings : Array[Array[Double]] = []
  for _ in 0.. 0.0 {
          let terms : Array[Double] = []
          let log_kernel = barycenter_cost_matrix(log_kernels, k)
          for i in 0.. log_barycenter_residual {
        log_barycenter_residual = change
      }
    }
    log_barycenter = next.log_barycenter
    iterations = iteration
    if log_barycenter_residual <= options.tolerance {
      termination = BarycenterTermination::Converged
      break
    }
  }

  match barycenter_materialize(log_barycenter, input.total_mass) {
    Err(_) => Err(BarycenterComputationError::NumericalFailure)
    Ok(barycenter) =>
      barycenter_evaluate_objective(
        barycenter,
        input.measures,
        input.scaled_weights,
        costs,
        options,
        log_barycenter_residual,
        iterations,
        termination,
      )
  }
}

///|
fn barycenter_cost_matrix(
  matrices : BarycenterCostMatrices,
  measure_index : Int,
) -> Array[Array[Double]] {
  match matrices {
    BarycenterCostMatrices::Shared(matrix) => matrix
    BarycenterCostMatrices::PerMeasure(per_measure) =>
      per_measure[measure_index]
  }
}

///|
fn barycenter_build_log_kernels(
  costs : BarycenterCostMatrices,
  epsilon : Double,
) -> Result[BarycenterCostMatrices, BarycenterComputationError] {
  match costs {
    BarycenterCostMatrices::Shared(cost) =>
      match barycenter_build_log_kernel(cost, epsilon) {
        Err(error) => Err(error)
        Ok(kernel) => Ok(BarycenterCostMatrices::Shared(kernel))
      }
    BarycenterCostMatrices::PerMeasure(per_measure) => {
      let kernels : Array[Array[Array[Double]]] = []
      for measure_index in 0.. return Err(error)
          Ok(kernel) => kernels.push(kernel)
        }
      }
      Ok(BarycenterCostMatrices::PerMeasure(kernels))
    }
  }
}

///|
fn barycenter_build_log_kernel(
  cost : Array[Array[Double]],
  epsilon : Double,
) -> Result[Array[Array[Double]], BarycenterComputationError] {
  let support_size = cost.length()
  let mut minimum_cost = cost[0][0]
  for i in 0.. BarycenterScalingState {
  let support_size = match log_kernels {
    BarycenterCostMatrices::Shared(matrix) => matrix.length()
    BarycenterCostMatrices::PerMeasure(matrices) => matrices[0].length()
  }
  let kernel_products : Array[Array[Double]] = []
  for _ in 0.. 0.0 {
          terms.push(log_kernel[i][j] + log_scalings[k][j])
        }
      }
      kernel_products[k][i] = logsumexp(terms)
    }
  }

  let weighted_logs = Array::make(support_size, 0.0)
  for i in 0.. Result[Array[Double], BarycenterError] {
  let barycenter = Array::make(log_barycenter.length(), 0.0)
  let mut actual_total = 0.0
  for i in 0.. Result[BarycenterResult, BarycenterComputationError] {
  let mut final_termination = termination
  let mut objective = 0.0
  let mut maximum_residual = 0.0
  for k in 0..
        return Err(BarycenterComputationError::SolverError(k, error))
      Ok(solution) => {
        objective = objective + weights[k] * solution.regularized_objective
        if !is_finite(objective) {
          return Err(BarycenterComputationError::NumericalFailure)
        }
        if solution.marginal_residual > maximum_residual {
          maximum_residual = solution.marginal_residual
        }
        if solution.termination is SinkhornTermination::IterationLimit {
          final_termination = BarycenterTermination::IterationLimit
        }
      }
    }
  }
  Ok({
    barycenter,
    weighted_objective: objective,
    maximum_transport_residual: maximum_residual,
    log_barycenter_residual: log_residual,
    iterations,
    termination: final_termination,
  })
}