///|
/// A scalar observation for batch regression.  `features` are copied so a
/// caller can reuse its input buffer after submitting the observation.
pub struct BatchObservation {
  features : Array[Double]
  value : Double
  weight : Double
  timestamp : Int
} derive(Debug)

///|
/// Construct a batch observation.
pub fn BatchObservation::new(
  features : Array[Double],
  value : Double,
  weight : Double,
  timestamp : Int,
) -> BatchObservation {
  {
    features: features.copy(),
    value,
    weight: if weight < 0.0 || weight.is_nan() {
      0.0
    } else {
      weight
    },
    timestamp,
  }
}

///|
/// Return the feature vector.
pub fn BatchObservation::features(self : BatchObservation) -> Array[Double] {
  self.features.copy()
}

///|
/// Return the observed value.
pub fn BatchObservation::value(self : BatchObservation) -> Double {
  self.value
}

///|
/// Return observation weight.
pub fn BatchObservation::weight(self : BatchObservation) -> Double {
  self.weight
}

///|
/// Return observation timestamp.
pub fn BatchObservation::timestamp(self : BatchObservation) -> Int {
  self.timestamp
}

///|
/// Return whether the observation is finite and usable.
pub fn BatchObservation::is_valid(self : BatchObservation) -> Bool {
  self.features.length() > 0 &&
  vector_is_finite(self.features) &&
  !self.value.is_nan() &&
  !self.value.is_inf() &&
  self.weight > 0.0
}

///|
/// A result returned by batch least-squares routines.
pub struct BatchFitResult {
  coefficients : Array[Double]
  covariance : Matrix
  residuals : Array[Double]
  weights : Array[Double]
  rank : Int
  rmse : Double
  mae : Double
  r_squared : Double
  condition : Double
  iterations : Int
  success : Bool
} derive(Debug)

///|
/// Return fitted coefficients.
pub fn BatchFitResult::coefficients(self : BatchFitResult) -> Array[Double] {
  self.coefficients.copy()
}

///|
/// Return coefficient covariance.
pub fn BatchFitResult::covariance(self : BatchFitResult) -> Matrix {
  self.covariance.copy()
}

///|
/// Return residuals.
pub fn BatchFitResult::residuals(self : BatchFitResult) -> Array[Double] {
  self.residuals.copy()
}

///|
/// Return final observation weights.
pub fn BatchFitResult::weights(self : BatchFitResult) -> Array[Double] {
  self.weights.copy()
}

///|
/// Return matrix rank.
pub fn BatchFitResult::rank(self : BatchFitResult) -> Int {
  self.rank
}

///|
/// Return root mean square error.
pub fn BatchFitResult::rmse(self : BatchFitResult) -> Double {
  self.rmse
}

///|
/// Return mean absolute error.
pub fn BatchFitResult::mae(self : BatchFitResult) -> Double {
  self.mae
}

///|
/// Return coefficient of determination.
pub fn BatchFitResult::r_squared(self : BatchFitResult) -> Double {
  self.r_squared
}

///|
/// Return design matrix condition estimate.
pub fn BatchFitResult::condition(self : BatchFitResult) -> Double {
  self.condition
}

///|
/// Return the number of solver iterations.
pub fn BatchFitResult::iterations(self : BatchFitResult) -> Int {
  self.iterations
}

///|
/// Return whether fitting succeeded.
pub fn BatchFitResult::success(self : BatchFitResult) -> Bool {
  self.success
}

///|
/// A regularized polynomial basis description.
pub struct PolynomialBasis {
  degree : Int
  center : Double
  scale : Double
} derive(Debug)

///|
/// Construct a polynomial basis with a numerically stable center and scale.
pub fn PolynomialBasis::new(
  degree : Int,
  center : Double,
  scale : Double,
) -> PolynomialBasis {
  {
    degree: if degree < 0 {
      0
    } else {
      degree
    },
    center,
    scale: if scale.abs() < 0.000000000001 {
      1.0
    } else {
      scale
    },
  }
}

///|
/// Return polynomial degree.
pub fn PolynomialBasis::degree(self : PolynomialBasis) -> Int {
  self.degree
}

///|
/// Return center.
pub fn PolynomialBasis::center(self : PolynomialBasis) -> Double {
  self.center
}

///|
/// Return scale.
pub fn PolynomialBasis::scale(self : PolynomialBasis) -> Double {
  self.scale
}

///|
/// Evaluate the polynomial feature vector at a scalar input.
pub fn PolynomialBasis::features(
  self : PolynomialBasis,
  x : Double,
) -> Array[Double] {
  let normalized = (x - self.center) / self.scale
  let result = Array::make(self.degree + 1, 1.0)
  for i in 1.. BatchFitResult {
  let count = if parameter_count < 0 { 0 } else { parameter_count }
  {
    coefficients: Array::make(count, 0.0),
    covariance: Matrix::zeros(count, count),
    residuals: [],
    weights: [],
    rank: 0,
    rmse: 0.0,
    mae: 0.0,
    r_squared: 0.0,
    condition: 0.0,
    iterations: 0,
    success: false,
  }
}

///|
/// Return the number of valid observations.
pub fn valid_batch_observation_count(
  observations : Array[BatchObservation],
) -> Int {
  let mut count = 0
  for observation in observations {
    if observation.is_valid() {
      count = count + 1
    }
  }
  count
}

///|
/// Infer the feature dimension from the first valid observation.
pub fn batch_feature_dimension(observations : Array[BatchObservation]) -> Int {
  for observation in observations {
    if observation.is_valid() {
      return observation.features().length()
    }
  }
  0
}

///|
/// Build a design matrix from observations, retaining only consistent rows.
pub fn batch_design_matrix(
  observations : Array[BatchObservation],
) -> (Matrix, Array[Double], Array[Double]) {
  let dimension = batch_feature_dimension(observations)
  let rows : Array[Array[Double]] = []
  let values : Array[Double] = []
  let weights : Array[Double] = []
  for observation in observations {
    let features = observation.features()
    if observation.is_valid() && features.length() == dimension {
      rows.push(features)
      values.push(observation.value())
      weights.push(observation.weight())
    }
  }
  (Matrix::from_rows(rows), values, weights)
}

///|
/// Form weighted normal equations `X'WX` and `X'Wy`.
pub fn batch_normal_equations(
  design : Matrix,
  values : Array[Double],
  weights : Array[Double],
) -> (Matrix, Array[Double]) {
  let parameters = design.cols()
  let normal = Matrix::zeros(parameters, parameters)
  let rhs = Array::make(parameters, 0.0)
  let count = if design.rows() < values.length() {
    design.rows()
  } else {
    values.length()
  }
  for i in 0.. 0.0 {
      let row = design.row(i)
      for j in 0.. ignore
        }
      }
    }
  }
  (normal, rhs)
}

///|
/// Solve a weighted linear least-squares problem with diagonal regularization.
pub fn batch_weighted_least_squares(
  observations : Array[BatchObservation],
  regularization : Double,
) -> BatchFitResult {
  let (design, values, weights) = batch_design_matrix(observations)
  if design.rows() == 0 || design.cols() == 0 {
    return empty_batch_fit(design.cols())
  }
  let (normal, rhs) = batch_normal_equations(design, values, weights)
  let regularized = normal.add_diagonal(
    if regularization < 0.0 {
      0.0
    } else {
      regularization
    },
  )
  let coefficients = match regularized.solve(rhs) {
    Solved(result) => result
    Singular | InvalidShape => return empty_batch_fit(design.cols())
  }
  let residuals = batch_residuals(design, values, coefficients)
  let metrics = batch_error_metrics(
    values,
    batch_predictions(design, coefficients),
  )
  let rank = design.rank(0.000000001)
  let covariance = match regularized.inverse() {
    Some(result) => result
    None => Matrix::zeros(design.cols(), design.cols())
  }
  {
    coefficients,
    covariance,
    residuals,
    weights,
    rank,
    rmse: metrics.0,
    mae: metrics.1,
    r_squared: batch_r_squared(values, batch_predictions(design, coefficients)),
    condition: regularized.condition_estimate(),
    iterations: 1,
    success: true,
  }
}

///|
/// Compute predictions from a design matrix and coefficient vector.
pub fn batch_predictions(
  design : Matrix,
  coefficients : Array[Double],
) -> Array[Double] {
  Array::makei(design.rows(), i => {
    let row = design.row(i)
    let count = if row.length() < coefficients.length() {
      row.length()
    } else {
      coefficients.length()
    }
    let mut value = 0.0
    for j in 0.. Array[Double] {
  let predictions = batch_predictions(design, coefficients)
  let count = if predictions.length() < values.length() {
    predictions.length()
  } else {
    values.length()
  }
  Array::makei(count, i => values[i] - predictions[i])
}

///|
/// Compute RMSE and MAE.
pub fn batch_error_metrics(
  actual : Array[Double],
  predicted : Array[Double],
) -> (Double, Double) {
  let count = if actual.length() < predicted.length() {
    actual.length()
  } else {
    predicted.length()
  }
  if count == 0 {
    return (0.0, 0.0)
  }
  let mut square = 0.0
  let mut absolute = 0.0
  for i in 0.. Double {
  let count = if actual.length() < predicted.length() {
    actual.length()
  } else {
    predicted.length()
  }
  if count == 0 {
    return 0.0
  }
  let mut mean = 0.0
  for i in 0.. BatchFitResult {
  let observations : Array[BatchObservation] = []
  let count = if values.length() < weights.length() {
    values.length()
  } else {
    weights.length()
  }
  for i in 0.. BatchFitResult {
  let observations : Array[BatchObservation] = []
  let count = if x.length() < y.length() { x.length() } else { y.length() }
  for i in 0.. Array[BatchObservation] {
  let result : Array[BatchObservation] = []
  let count = if x.length() < y.length() { x.length() } else { y.length() }
  for i in 0.. BatchFitResult {
  let basis = PolynomialBasis::new(degree, center, scale)
  batch_weighted_least_squares(
    polynomial_observations(x, y, basis),
    0.000000001,
  )
}

///|
/// Evaluate a fitted polynomial at x.
pub fn evaluate_polynomial(
  coefficients : Array[Double],
  x : Double,
  center : Double,
  scale : Double,
) -> Double {
  let safe_scale = if scale.abs() < 0.000000000001 { 1.0 } else { scale }
  let normalized = (x - center) / safe_scale
  let mut power = 1.0
  let mut result = 0.0
  for coefficient in coefficients {
    result = result + coefficient * power
    power = power * normalized
  }
  result
}

///|
/// Compute leverage values from a design matrix using a regularized inverse.
pub fn batch_leverage(
  design : Matrix,
  regularization : Double,
) -> Array[Double] {
  let gram = design
    .transpose()
    .multiply(design)
    .add_diagonal(if regularization < 0.0 { 0.0 } else { regularization })
  guard gram.inverse() is Some(inverse) else {
    return Array::make(design.rows(), 0.0)
  }
  let result = Array::make(design.rows(), 0.0)
  for i in 0.. Array[Double] {
  let count = if residuals.length() < leverage.length() {
    residuals.length()
  } else {
    leverage.length()
  }
  let safe_mse = if mse <= 0.0 { 1.0 } else { mse }
  Array::makei(count, i => {
    residuals[i] *
    residuals[i] /
    safe_mse *
    leverage[i] /
    (1.0 - leverage[i]).clamp(min=0.000001, max=1.0)
  })
}

///|
/// A recursive least-squares estimator for slowly changing calibration.
pub struct RecursiveLeastSquares {
  coefficients : Array[Double]
  mut covariance : Matrix
  forgetting : Double
  mut updates : Int
  mut rejected : Int
} derive(Debug)

///|
/// Construct a recursive least-squares estimator.
pub fn RecursiveLeastSquares::new(
  dimension : Int,
  initial_covariance : Double,
  forgetting : Double,
) -> RecursiveLeastSquares {
  let size = if dimension < 0 { 0 } else { dimension }
  {
    coefficients: Array::make(size, 0.0),
    covariance: Matrix::diagonal(
      size,
      if initial_covariance <= 0.0 {
        1.0
      } else {
        initial_covariance
      },
    ),
    forgetting: forgetting.clamp(min=0.000001, max=1.0),
    updates: 0,
    rejected: 0,
  }
}

///|
/// Return current coefficients.
pub fn RecursiveLeastSquares::coefficients(
  self : RecursiveLeastSquares,
) -> Array[Double] {
  self.coefficients.copy()
}

///|
/// Return parameter covariance.
pub fn RecursiveLeastSquares::covariance(
  self : RecursiveLeastSquares,
) -> Matrix {
  self.covariance.copy()
}

///|
/// Return forgetting factor.
pub fn RecursiveLeastSquares::forgetting(
  self : RecursiveLeastSquares,
) -> Double {
  self.forgetting
}

///|
/// Return update count.
pub fn RecursiveLeastSquares::updates(self : RecursiveLeastSquares) -> Int {
  self.updates
}

///|
/// Return rejected update count.
pub fn RecursiveLeastSquares::rejected(self : RecursiveLeastSquares) -> Int {
  self.rejected
}

///|
/// Predict one scalar response.
pub fn RecursiveLeastSquares::predict(
  self : RecursiveLeastSquares,
  features : Array[Double],
) -> Double {
  let count = if features.length() < self.coefficients.length() {
    features.length()
  } else {
    self.coefficients.length()
  }
  let mut result = 0.0
  for i in 0.. Double? {
  if features.length() != self.coefficients.length() ||
    features.length() == 0 ||
    value.is_nan() ||
    value.is_inf() {
    self.rejected = self.rejected + 1
    return None
  }
  let prediction = self.predict(features)
  let innovation = value - prediction
  let prior = self.covariance.multiply_vector(features)
  let mut denominator = self.forgetting
  for i in 0.. prior[i] / denominator)
  for i in 0.. ignore
    }
  }
  self.covariance = next.symmetric_part()
  self.updates = self.updates + 1
  Some(innovation)
}

///|
/// Reset parameters to zero with a supplied covariance scale.
pub fn RecursiveLeastSquares::reset(
  self : RecursiveLeastSquares,
  covariance_scale : Double,
) -> Unit {
  for i in 0.. RecursiveLeastSquares {
  let estimator = RecursiveLeastSquares::new(
    batch_feature_dimension(observations),
    initial_covariance,
    forgetting,
  )
  for observation in observations {
    if observation.is_valid() {
      estimator.update(observation.features(), observation.value()) |> ignore
    }
  }
  estimator
}

///|
/// Compute a weighted residual median for regression diagnostics.
pub fn batch_residual_median(
  residuals : Array[Double],
  weights : Array[Double],
) -> Double {
  robust_weighted_quantile(residuals, weights, 0.5)
}

///|
/// Return a robust batch fit by iteratively updating observation weights.
pub fn batch_robust_fit(
  observations : Array[BatchObservation],
  config : RobustEstimatorConfig,
) -> BatchFitResult {
  let initial = batch_weighted_least_squares(observations, 0.000000001)
  if !initial.success() {
    return initial
  }
  let mut mutable_weights = initial.weights()
  let mut current = initial
  for _ in 0.. Double {
  let n = fit.residuals().length().to_double()
  let k = fit.coefficients().length().to_double()
  if n <= 0.0 {
    0.0
  } else {
    n * (fit.rmse() * fit.rmse()).max(0.000000000001) + 2.0 * k
  }
}

///|
/// Compute BIC for a batch fit.
pub fn batch_bic(fit : BatchFitResult) -> Double {
  let n = fit.residuals().length().to_double()
  let k = fit.coefficients().length().to_double()
  if n <= 0.0 {
    0.0
  } else {
    n * (fit.rmse() * fit.rmse()).max(0.000000000001) + k * n.max(1.0).sqrt()
  }
}

///|
/// Return a compact report line for a batch fit.
pub fn batch_fit_summary(fit : BatchFitResult) -> String {
  "rank=" +
  fit.rank().to_string() +
  ", rmse=" +
  fit.rmse().to_string() +
  ", mae=" +
  fit.mae().to_string() +
  ", r2=" +
  fit.r_squared().to_string() +
  ", condition=" +
  fit.condition().to_string()
}

///|
/// Build observations for a scalar time series with an intercept and time.
pub fn time_series_observations(
  timestamps : Array[Int],
  values : Array[Double],
) -> Array[BatchObservation] {
  let result : Array[BatchObservation] = []
  let count = if timestamps.length() < values.length() {
    timestamps.length()
  } else {
    values.length()
  }
  if count == 0 {
    return result
  }
  let origin = timestamps[0]
  for i in 0.. BatchFitResult {
  batch_weighted_least_squares(
    time_series_observations(timestamps, values),
    0.000000001,
  )
}

///|
/// Predict values on a timestamp grid from a time-series trend fit.
pub fn predict_time_series_trend(
  fit : BatchFitResult,
  origin : Int,
  timestamps : Array[Int],
) -> Array[Double] {
  let coefficients = fit.coefficients()
  Array::makei(timestamps.length(), i => {
    if coefficients.length() < 2 {
      0.0
    } else {
      coefficients[0] + coefficients[1] * (timestamps[i] - origin).to_double()
    }
  })
}

///|
/// Compute a prediction interval radius from fit error and confidence scale.
pub fn batch_prediction_radius(
  fit : BatchFitResult,
  confidence_scale : Double,
) -> Double {
  let multiplier = if confidence_scale < 0.0 { 0.0 } else { confidence_scale }
  fit.rmse() * multiplier
}

///|
/// Return an interval around a predicted value.
pub fn batch_prediction_interval(
  fit : BatchFitResult,
  prediction : Double,
  confidence_scale : Double,
) -> (Double, Double) {
  let radius = batch_prediction_radius(fit, confidence_scale)
  (prediction - radius, prediction + radius)
}

///|
/// Return whether a batch fit meets numerical quality limits.
pub fn batch_fit_is_usable(
  fit : BatchFitResult,
  maximum_condition : Double,
  maximum_rmse : Double,
) -> Bool {
  fit.success() &&
  fit.rank() > 0 &&
  fit.condition() <= maximum_condition &&
  fit.rmse() <= maximum_rmse
}