///|
pub struct LinearRegressionResult {
  slope : Double
  intercept : Double
  scale : Double
  r_squared : Double
  iterations : Int
  converged : Bool
  residuals : Array[Double]
}

///|
pub fn linear_regression(
  x : Array[Double],
  y : Array[Double],
) -> LinearRegressionResult {
  if x.length() == 0 || x.length() != y.length() {
    return {
      slope: 0.0,
      intercept: 0.0,
      scale: 0.0,
      r_squared: 0.0,
      iterations: 0,
      converged: false,
      residuals: [],
    }
  }
  let center_x = mean(x)
  let center_y = mean(y)
  let mut numerator = 0.0
  let mut denominator = 0.0
  for index = 0; index < x.length(); index = index + 1 {
    numerator += (x[index] - center_x) * (y[index] - center_y)
    let delta = x[index] - center_x
    denominator += delta * delta
  }
  let slope = if denominator == 0.0 { 0.0 } else { numerator / denominator }
  let intercept = center_y - slope * center_x
  let predictions = predict_linear(x, slope, intercept)
  let residuals = []
  for index = 0; index < y.length(); index = index + 1 {
    residuals.push(y[index] - predictions[index])
  }
  {
    slope,
    intercept,
    scale: mad(residuals),
    r_squared: r_squared(y, predictions),
    iterations: 1,
    converged: true,
    residuals,
  }
}

///|
pub fn predict_linear(
  x : Array[Double],
  slope : Double,
  intercept : Double,
) -> Array[Double] {
  let result = []
  for value in x {
    result.push(intercept + slope * value)
  }
  result
}

///|
pub fn huber_regression(
  x : Array[Double],
  y : Array[Double],
  tuning? : Double = 1.345,
  max_iter? : Int = 50,
  tol? : Double = 0.0001,
) -> LinearRegressionResult {
  if x.length() == 0 || x.length() != y.length() {
    return linear_regression(x, y)
  }
  if tuning <= 0.0 || max_iter < 1 || tol <= 0.0 {
    abort("invalid robust regression configuration")
  }
  let ordinary = linear_regression(x, y)
  let mut slope = ordinary.slope
  let mut intercept = ordinary.intercept
  let mut iterations = 0
  let mut converged = false
  let initial_scale = if ordinary.scale == 0.0 {
    sample_stddev(y)
  } else {
    ordinary.scale
  }
  let scale = if initial_scale == 0.0 { 1.0 } else { initial_scale }
  for iteration = 0; iteration < max_iter; iteration = iteration + 1 {
    let residuals = []
    for index = 0; index < x.length(); index = index + 1 {
      residuals.push(y[index] - (intercept + slope * x[index]))
    }
    let cutoff = tuning * scale
    let weights = []
    for residual in residuals {
      weights.push(huber_weight(residual, cutoff))
    }
    let weighted_x = weighted_mean(x, weights)
    let weighted_y = weighted_mean(y, weights)
    let mut numerator = 0.0
    let mut denominator = 0.0
    for index = 0; index < x.length(); index = index + 1 {
      let dx = x[index] - weighted_x
      numerator += weights[index] * dx * (y[index] - weighted_y)
      denominator += weights[index] * dx * dx
    }
    let next_slope = if denominator == 0.0 {
      slope
    } else {
      numerator / denominator
    }
    let next_intercept = weighted_y - next_slope * weighted_x
    iterations = iteration + 1
    if abs_double(next_slope - slope) <= tol &&
      abs_double(next_intercept - intercept) <= tol {
      slope = next_slope
      intercept = next_intercept
      converged = true
      break
    }
    slope = next_slope
    intercept = next_intercept
  }
  let predictions = predict_linear(x, slope, intercept)
  let residuals = []
  for index = 0; index < y.length(); index = index + 1 {
    residuals.push(y[index] - predictions[index])
  }
  {
    slope,
    intercept,
    scale: mad(residuals),
    r_squared: r_squared(y, predictions),
    iterations,
    converged,
    residuals,
  }
}

///|
pub fn weighted_linear_regression(
  x : Array[Double],
  y : Array[Double],
  weights : Array[Double],
) -> LinearRegressionResult {
  if x.length() == 0 ||
    x.length() != y.length() ||
    x.length() != weights.length() {
    return linear_regression(x, y)
  }
  let center_x = weighted_mean(x, weights)
  let center_y = weighted_mean(y, weights)
  let mut numerator = 0.0
  let mut denominator = 0.0
  for index = 0; index < x.length(); index = index + 1 {
    let dx = x[index] - center_x
    numerator += weights[index] * dx * (y[index] - center_y)
    denominator += weights[index] * dx * dx
  }
  let slope = if denominator == 0.0 { 0.0 } else { numerator / denominator }
  let intercept = center_y - slope * center_x
  let predictions = predict_linear(x, slope, intercept)
  let residuals = []
  for index = 0; index < y.length(); index = index + 1 {
    residuals.push(y[index] - predictions[index])
  }
  {
    slope,
    intercept,
    scale: weighted_variance(residuals, weights).sqrt(),
    r_squared: r_squared(y, predictions),
    iterations: 1,
    converged: true,
    residuals,
  }
}

///|
pub fn theil_sen_slope(x : Array[Double], y : Array[Double]) -> Double {
  if x.length() < 2 || x.length() != y.length() {
    return 0.0
  }
  let slopes = []
  for left = 0; left < x.length(); left = left + 1 {
    for right = left + 1; right < x.length(); right = right + 1 {
      if x[right] != x[left] {
        slopes.push((y[right] - y[left]) / (x[right] - x[left]))
      }
    }
  }
  median(slopes)
}

///|
pub fn theil_sen_regression(
  x : Array[Double],
  y : Array[Double],
) -> LinearRegressionResult {
  if x.length() == 0 || x.length() != y.length() {
    return linear_regression(x, y)
  }
  let slope = theil_sen_slope(x, y)
  let intercepts = []
  for index = 0; index < x.length(); index = index + 1 {
    intercepts.push(y[index] - slope * x[index])
  }
  let intercept = median(intercepts)
  let predictions = predict_linear(x, slope, intercept)
  let residuals = []
  for index = 0; index < y.length(); index = index + 1 {
    residuals.push(y[index] - predictions[index])
  }
  {
    slope,
    intercept,
    scale: mad(residuals),
    r_squared: r_squared(y, predictions),
    iterations: 1,
    converged: true,
    residuals,
  }
}

///|
pub fn regression_prediction_interval(
  model : LinearRegressionResult,
  x_value : Double,
  z_value : Double,
) -> Array[Double] {
  let prediction = model.intercept + model.slope * x_value
  [prediction - z_value * model.scale, prediction + z_value * model.scale]
}

///|
pub fn regression_residual_scale(model : LinearRegressionResult) -> Double {
  mad(model.residuals)
}

///|
pub fn regression_leverage(x : Array[Double], x_value : Double) -> Double {
  if x.length() == 0 {
    return 0.0
  }
  let center = mean(x)
  let mut total = 0.0
  for value in x {
    let delta = value - center
    total += delta * delta
  }
  if total == 0.0 {
    1.0 / x.length().to_double()
  } else {
    1.0 / x.length().to_double() +
    (x_value - center) * (x_value - center) / total
  }
}

///|
pub fn robust_regression_weights(
  model : LinearRegressionResult,
  tuning : Double,
) -> Array[Double] {
  let result = []
  for residual in model.residuals {
    result.push(
      huber_weight(
        residual,
        if model.scale == 0.0 {
          tuning
        } else {
          tuning * model.scale
        },
      ),
    )
  }
  result
}

///|
pub fn residual_outlier_indices(
  model : LinearRegressionResult,
  threshold? : Double = 3.5,
) -> Array[Int] {
  outlier_indices_z(model.residuals, threshold~)
}

///|
pub fn regression_mae(model : LinearRegressionResult) -> Double {
  let denominator = if model.residuals.length() == 0 {
    1.0
  } else {
    model.residuals.length().to_double()
  }
  sum_absolute(model.residuals) / denominator
}

///|
pub fn regression_rmse(model : LinearRegressionResult) -> Double {
  root_mean_square(model.residuals)
}

///|
pub fn median_regression_error(model : LinearRegressionResult) -> Double {
  median(model.residuals)
}