///|
/// A measured point with an explicit run identifier and optional uncertainty.
pub(all) struct MeasuredPoint {
  run_id : String
  x : Double
  y : Double
  standard_uncertainty : Double
  replicate : Int
} derive(Debug, Eq, ToJson)

///|
pub fn MeasuredPoint::new(
  run_id~ : String,
  x~ : Double,
  y~ : Double,
  standard_uncertainty? : Double = 0.0,
  replicate? : Int = 1,
) -> MeasuredPoint raise @domain.ExtractionError {
  if run_id.trim().length() == 0 {
    raise @domain.ExtractionError::InvalidArgument(
      message="measurement run id cannot be blank",
    )
  }
  @domain.require_finite("measurement x", x)
  @domain.require_finite("measurement y", y)
  @domain.require_non_negative("measurement uncertainty", standard_uncertainty)
  if replicate < 1 {
    raise @domain.ExtractionError::InvalidArgument(
      message="replicate number must be positive",
    )
  }
  { run_id, x, y, standard_uncertainty, replicate }
}

///|
pub fn MeasuredPoint::weight(self : MeasuredPoint) -> Double {
  if self.standard_uncertainty <= 0.0 {
    1.0
  } else {
    1.0 / (self.standard_uncertainty * self.standard_uncertainty)
  }
}

///|
pub(all) struct ResidualRecord {
  run_id : String
  x : Double
  observed : Double
  predicted : Double
  residual : Double
  standardized_residual : Double
  leverage : Double
} derive(Debug, Eq, ToJson)

///|
pub(all) struct RegressionDiagnostics {
  count : Int
  residuals : Array[ResidualRecord]
  mean_absolute_error : Double
  median_absolute_error : Double
  root_mean_squared_error : Double
  maximum_absolute_error : Double
  bias : Double
  r_squared : Double
  outlier_count : Int
} derive(Debug, Eq, ToJson)

///|
fn median(values : ArrayView[Double]) -> Double raise @domain.ExtractionError {
  if values.length() == 0 {
    raise @domain.ExtractionError::InsufficientData(required=1, actual=0)
  }
  let sorted = values.to_owned()
  sorted.sort()
  let middle = sorted.length() / 2
  if sorted.length() % 2 == 0 {
    (sorted[middle - 1] + sorted[middle]) / 2.0
  } else {
    sorted[middle]
  }
}

///|
/// Build residual and leverage diagnostics for an existing polynomial model.
pub fn diagnose(
  model : RegressionModel,
  points : ArrayView[MeasuredPoint],
  outlier_threshold? : Double = 3.0,
) -> RegressionDiagnostics raise @domain.ExtractionError {
  if points.length() == 0 {
    raise @domain.ExtractionError::InsufficientData(required=1, actual=0)
  }
  @domain.require_positive("outlier threshold", outlier_threshold)
  let mut sum_abs = 0.0
  let mut sum_squared = 0.0
  let mut sum_error = 0.0
  let mut sum_y = 0.0
  for point in points {
    @domain.require_finite("measured x", point.x)
    @domain.require_finite("measured y", point.y)
    sum_y += point.y
  }
  let mean_y = sum_y / points.length().to_double()
  let mut total = 0.0
  let mut maximum = 0.0
  let raw_residuals : Array[Double] = []
  for point in points {
    let predicted = model.predict(point.x)
    let error = point.y - predicted
    raw_residuals.push(error)
    sum_abs += error.abs()
    sum_squared += error * error
    sum_error += error
    total += (point.y - mean_y) * (point.y - mean_y)
    maximum = maximum.max(error.abs())
  }
  let standard_error = (sum_squared / points.length().to_double())
    .sqrt()
    .max(1.0e-12)
  let residuals : Array[ResidualRecord] = []
  let mut outliers = 0
  for index, point in points {
    let predicted = model.predict(point.x)
    let error = point.y - predicted
    let standardized = error / standard_error
    let leverage = 1.0 / points.length().to_double() +
      (point.x - points[0].x) *
      (point.x - points[0].x) /
      (
        1.0 +
        points.fold(init=0.0, (sum, item) => {
          sum + (item.x - points[0].x) * (item.x - points[0].x)
        })
      )
    if standardized.abs() > outlier_threshold {
      outliers += 1
    }
    residuals.push({
      run_id: point.run_id,
      x: point.x,
      observed: point.y,
      predicted,
      residual: error,
      standardized_residual: standardized,
      leverage,
    })
    ignore(index)
  }
  let r2 = if total <= 1.0e-20 { 1.0 } else { 1.0 - sum_squared / total }
  {
    count: points.length(),
    residuals,
    mean_absolute_error: sum_abs / points.length().to_double(),
    median_absolute_error: median([ for value in raw_residuals => value.abs() ]),
    root_mean_squared_error: (sum_squared / points.length().to_double()).sqrt(),
    maximum_absolute_error: maximum,
    bias: sum_error / points.length().to_double(),
    r_squared: r2,
    outlier_count: outliers,
  }
}

///|
pub(all) struct PredictionInterval {
  x : Double
  prediction : Double
  lower : Double
  upper : Double
  half_width : Double
  confidence_level : Double
} derive(Debug, Eq, ToJson)

///|
/// Estimate a prediction interval using model residual dispersion.
pub fn prediction_interval(
  model : RegressionModel,
  points : ArrayView[MeasuredPoint],
  x : Double,
  confidence_level? : Double = 0.95,
) -> PredictionInterval raise @domain.ExtractionError {
  @domain.require_finite("prediction x", x)
  if confidence_level <= 0.0 || confidence_level >= 1.0 {
    raise @domain.ExtractionError::InvalidArgument(
      message="confidence level must be between zero and one",
    )
  }
  let diagnostics = diagnose(model, points)
  let z = if confidence_level >= 0.995 {
    2.807
  } else if confidence_level >= 0.975 {
    2.241
  } else if confidence_level >= 0.95 {
    1.96
  } else if confidence_level >= 0.90 {
    1.645
  } else {
    1.282
  }
  let width = z * diagnostics.root_mean_squared_error
  let prediction = model.predict(x)
  {
    x,
    prediction,
    lower: prediction - width,
    upper: prediction + width,
    half_width: width,
    confidence_level,
  }
}

///|
/// A robust fit result after iteratively down-weighting large residuals.
pub(all) struct RobustFit {
  model : RegressionModel
  iterations : Int
  removed_points : Int
  diagnostics : RegressionDiagnostics
} derive(Debug, Eq, ToJson)

///|
fn weighted_observations(
  points : ArrayView[MeasuredPoint],
  model : RegressionModel,
  cutoff : Double,
) -> Array[Observation] {
  [
    for point in points => {
      let error = (point.y - model.predict(point.x)).abs()
      let robust_weight = if error <= cutoff { 1.0 } else { cutoff / error }
      Observation::new(
        x=point.x,
        y=point.y,
        weight=point.weight() * robust_weight,
      ) catch {
        _ => Observation::{ x: point.x, y: point.y, weight: 1.0 }
      }
    }
  ]
}

///|
/// Iteratively reweighted least squares for noisy equilibrium measurements.
pub fn robust_polynomial_fit(
  points : ArrayView[MeasuredPoint],
  degree : Int,
  maximum_iterations? : Int = 8,
  cutoff? : Double = 2.5,
) -> RobustFit raise @domain.ExtractionError {
  if points.length() < degree + 1 {
    raise @domain.ExtractionError::InsufficientData(
      required=degree + 1,
      actual=points.length(),
    )
  }
  if maximum_iterations < 1 {
    raise @domain.ExtractionError::InvalidArgument(
      message="robust fit iteration limit must be positive",
    )
  }
  @domain.require_positive("robust fit cutoff", cutoff)
  let initial = [
    for point in points => {
      Observation::new(x=point.x, y=point.y, weight=point.weight())
    }
  ]
  let mut model = fit_polynomial(initial, degree)
  let mut removed = 0
  let mut iterations = 0
  for index in 0.. 1.0e-9 {
        changed = true
      }
    }
    model = next_model
    if !changed {
      break
    }
  }
  let diagnostics = diagnose(model, points, outlier_threshold=cutoff)
  removed = diagnostics.outlier_count
  { model, iterations, removed_points: removed, diagnostics }
}

///|
/// Evaluate a model on a regular grid for plotting or process screening.
pub fn prediction_grid(
  model : RegressionModel,
  lower~ : Double,
  upper~ : Double,
  count~ : Int,
) -> Array[PredictionInterval] raise @domain.ExtractionError {
  if count < 2 {
    raise @domain.ExtractionError::InsufficientData(required=2, actual=count)
  }
  @domain.require_finite("prediction grid lower", lower)
  @domain.require_finite("prediction grid upper", upper)
  if upper <= lower {
    raise @domain.ExtractionError::InvalidArgument(
      message="prediction grid bounds must be increasing",
    )
  }
  [
    for index in 0.. {
      let x = lower +
        (upper - lower) * index.to_double() / (count - 1).to_double()
      let prediction = model.predict(x)
      {
        x,
        prediction,
        lower: prediction,
        upper: prediction,
        half_width: 0.0,
        confidence_level: 1.0,
      }
    }
  ]
}

///|
/// Calculate a normalized sensitivity of a fitted correlation at one point.
pub fn model_sensitivity(
  model : RegressionModel,
  x : Double,
  relative_step? : Double = 1.0e-4,
) -> Double raise @domain.ExtractionError {
  @domain.require_finite("sensitivity x", x)
  @domain.require_positive("sensitivity step", relative_step)
  let step = relative_step * x.abs().max(1.0)
  let base = model.predict(x)
  let slope = (model.predict(x + step) - model.predict(x - step)) / (2.0 * step)
  if base.abs() <= 1.0e-12 {
    slope
  } else {
    slope * x / base
  }
}

///|
/// Compare two fitted models on the same measured points.
pub(all) struct ModelComparison {
  first_rmse : Double
  second_rmse : Double
  first_mae : Double
  second_mae : Double
  preferred : String
  improvement_fraction : Double
} derive(Debug, Eq, ToJson)

///|
pub fn compare_models(
  first : RegressionModel,
  second : RegressionModel,
  points : ArrayView[MeasuredPoint],
) -> ModelComparison raise @domain.ExtractionError {
  let first_diagnostics = diagnose(first, points)
  let second_diagnostics = diagnose(second, points)
  let preferred = if first_diagnostics.root_mean_squared_error <=
    second_diagnostics.root_mean_squared_error {
    "first"
  } else {
    "second"
  }
  let best = first_diagnostics.root_mean_squared_error.min(
    second_diagnostics.root_mean_squared_error,
  )
  let worst = first_diagnostics.root_mean_squared_error.max(
    second_diagnostics.root_mean_squared_error,
  )
  {
    first_rmse: first_diagnostics.root_mean_squared_error,
    second_rmse: second_diagnostics.root_mean_squared_error,
    first_mae: first_diagnostics.mean_absolute_error,
    second_mae: second_diagnostics.mean_absolute_error,
    preferred,
    improvement_fraction: if worst <= 1.0e-12 {
      0.0
    } else {
      (worst - best) / worst
    },
  }
}