///|
/// Missing-value imputation strategy.
pub enum ImputationStrategy {
  Forward
  Backward
  MedianFill
  Linear
  RollingMedian
  RollingMean
  Interpolate
}

///|
/// Description of one missing run.
pub struct MissingGap {
  start : Int
  end : Int
  length : Int
  left_value : Double
  right_value : Double
}

///|
/// Result of an imputation pass.
pub struct ImputationReport {
  original : Array[Double]
  filled : Array[Double]
  missing_before : Int
  missing_after : Int
  gaps : Array[MissingGap]
  changed : Int
  quality_before : Double
  quality_after : Double
}

///|
pub fn imputation_strategy_catalog() -> Array[ImputationStrategy] {
  [
    Forward,
    Backward,
    MedianFill,
    Linear,
    RollingMedian,
    RollingMean,
    Interpolate,
  ]
}

///|
pub fn imputation_missing_count(data : Array[Double]) -> Int {
  quality_missing_indices(data).length()
}

///|
pub fn imputation_gap_ranges(data : Array[Double]) -> Array[MissingGap] {
  let result = []
  let mut index = 0
  while index < data.length() {
    if !quality_is_missing(data[index]) {
      index += 1
    } else {
      let start = index
      while index < data.length() && quality_is_missing(data[index]) {
        index += 1
      }
      let end = index
      let left_value = if start == 0 { 0.0 } else { data[start - 1] }
      let right_value = if end >= data.length() {
        left_value
      } else {
        data[end]
      }
      result.push({ start, end, length: end - start, left_value, right_value })
    }
  }
  result
}

///|
pub fn imputation_gap_lengths(data : Array[Double]) -> Array[Int] {
  let result = []
  for gap in imputation_gap_ranges(data) {
    result.push(gap.length)
  }
  result
}

///|
pub fn imputation_max_gap(data : Array[Double]) -> Int {
  let mut result = 0
  for length in imputation_gap_lengths(data) {
    if length > result {
      result = length
    }
  }
  result
}

///|
pub fn imputation_forward(data : Array[Double]) -> Array[Double] {
  quality_impute_forward(data)
}

///|
pub fn imputation_backward(data : Array[Double]) -> Array[Double] {
  let reversed = []
  for index = data.length() - 1; index >= 0; index = index - 1 {
    reversed.push(data[index])
  }
  let filled = quality_impute_forward(reversed)
  let result = []
  for index = filled.length() - 1; index >= 0; index = index - 1 {
    result.push(filled[index])
  }
  result
}

///|
pub fn imputation_median(data : Array[Double]) -> Array[Double] {
  quality_impute_median(data)
}

///|
pub fn imputation_linear(data : Array[Double]) -> Array[Double] {
  let result = data.copy()
  let gaps = imputation_gap_ranges(data)
  for gap in gaps {
    for index = gap.start; index < gap.end; index = index + 1 {
      let weight = if gap.length == 0 {
        0.0
      } else {
        (index - gap.start + 1).to_double() / (gap.length + 1).to_double()
      }
      result[index] = gap.left_value +
        (gap.right_value - gap.left_value) * weight
    }
  }
  result
}

///|
pub fn imputation_interpolate_bounded(
  data : Array[Double],
  lower : Double,
  upper : Double,
) -> Array[Double] {
  clip_range(imputation_linear(data), lower, upper)
}

///|
pub fn imputation_rolling_median(
  data : Array[Double],
  window : Int,
) -> Array[Double] {
  let result = data.copy()
  let baseline = rolling_median(
    imputation_median(data),
    if window < 2 {
      2
    } else {
      window
    },
  )
  for index = 0; index < result.length(); index = index + 1 {
    if quality_is_missing(result[index]) {
      result[index] = baseline[index]
    }
  }
  result
}

///|
pub fn imputation_rolling_mean(
  data : Array[Double],
  window : Int,
) -> Array[Double] {
  let result = data.copy()
  let baseline = rolling_mean(
    imputation_median(data),
    if window < 2 {
      2
    } else {
      window
    },
  )
  for index = 0; index < result.length(); index = index + 1 {
    if quality_is_missing(result[index]) {
      result[index] = baseline[index]
    }
  }
  result
}

///|
pub fn imputation_with_strategy(
  data : Array[Double],
  strategy : ImputationStrategy,
  window : Int,
) -> Array[Double] {
  match strategy {
    Forward => imputation_forward(data)
    Backward => imputation_backward(data)
    MedianFill => imputation_median(data)
    Linear => imputation_linear(data)
    RollingMedian => imputation_rolling_median(data, window)
    RollingMean => imputation_rolling_mean(data, window)
    Interpolate =>
      imputation_interpolate_bounded(
        data,
        min_value(imputation_median(data)),
        max_value(imputation_median(data)),
      )
  }
}

///|
pub fn imputation_changed_count(
  original : Array[Double],
  filled : Array[Double],
) -> Int {
  let limit = if original.length() < filled.length() {
    original.length()
  } else {
    filled.length()
  }
  let mut count = 0
  for index = 0; index < limit; index = index + 1 {
    if original[index] != filled[index] {
      count += 1
    }
  }
  count
}

///|
pub fn imputation_report(
  data : Array[Double],
  strategy : ImputationStrategy,
  window : Int,
  rule : QualityRule,
) -> ImputationReport {
  let filled = imputation_with_strategy(data, strategy, window)
  {
    original: data.copy(),
    filled,
    missing_before: imputation_missing_count(data),
    missing_after: imputation_missing_count(filled),
    gaps: imputation_gap_ranges(data),
    changed: imputation_changed_count(data, filled),
    quality_before: quality_report(data, rule).quality_score,
    quality_after: quality_report(filled, rule).quality_score,
  }
}

///|
pub fn imputation_report_vector(report : ImputationReport) -> Array[Double] {
  [
    report.missing_before.to_double(),
    report.missing_after.to_double(),
    report.gaps.length().to_double(),
    imputation_max_gap(report.original).to_double(),
    report.changed.to_double(),
    report.quality_before,
    report.quality_after,
  ]
}

///|
pub fn imputation_report_lines(report : ImputationReport) -> Array[String] {
  [
    "missing_before=" + report.missing_before.to_string(),
    "missing_after=" + report.missing_after.to_string(),
    "gap_count=" + report.gaps.length().to_string(),
    "changed=" + report.changed.to_string(),
    "quality_before=" + report.quality_before.to_string(),
    "quality_after=" + report.quality_after.to_string(),
  ]
}

///|
pub fn imputation_report_string(report : ImputationReport) -> String {
  imputation_report_lines(report).join("\n")
}

///|
pub fn imputation_quality_gain(report : ImputationReport) -> Double {
  report.quality_after - report.quality_before
}

///|
pub fn imputation_method_scores(
  data : Array[Double],
  window : Int,
  rule : QualityRule,
) -> Array[Double] {
  let result = []
  for strategy in imputation_strategy_catalog() {
    result.push(
      imputation_quality_gain(imputation_report(data, strategy, window, rule)),
    )
  }
  result
}

///|
pub fn imputation_best_strategy(
  data : Array[Double],
  window : Int,
  rule : QualityRule,
) -> ImputationStrategy {
  let strategies = imputation_strategy_catalog()
  if strategies.length() == 0 {
    return MedianFill
  }
  let mut best = strategies[0]
  let mut score = imputation_quality_gain(
    imputation_report(data, best, window, rule),
  )
  for strategy in strategies {
    let current = imputation_quality_gain(
      imputation_report(data, strategy, window, rule),
    )
    if current > score {
      best = strategy
      score = current
    }
  }
  best
}

///|
pub fn imputation_consensus(
  data : Array[Double],
  window : Int,
) -> Array[Double] {
  let candidates = []
  for strategy in imputation_strategy_catalog() {
    candidates.push(imputation_with_strategy(data, strategy, window))
  }
  let result = []
  for index = 0; index < data.length(); index = index + 1 {
    if !quality_is_missing(data[index]) {
      result.push(data[index])
    } else {
      let values = []
      for candidate in candidates {
        values.push(candidate[index])
      }
      result.push(median(values))
    }
  }
  result
}

///|
pub fn imputation_sensitivity(
  data : Array[Double],
  strategy : ImputationStrategy,
  windows : Array[Int],
) -> Array[Double] {
  let result = []
  for window in windows {
    let filled = imputation_with_strategy(data, strategy, window)
    result.push(mad(filled))
  }
  result
}

///|
pub fn imputation_gap_fill_values(
  data : Array[Double],
  strategy : ImputationStrategy,
  window : Int,
) -> Array[Double] {
  let filled = imputation_with_strategy(data, strategy, window)
  let result = []
  for gap in imputation_gap_ranges(data) {
    for index = gap.start; index < gap.end; index = index + 1 {
      result.push(filled[index])
    }
  }
  result
}

///|
pub fn imputation_mask(data : Array[Double]) -> Array[Bool] {
  let result = []
  for value in data {
    result.push(quality_is_missing(value))
  }
  result
}

///|
pub fn imputation_restore(
  filled : Array[Double],
  original : Array[Double],
  mask : Array[Bool],
) -> Array[Double] {
  let result = filled.copy()
  let limit = if result.length() < original.length() {
    result.length()
  } else {
    original.length()
  }
  for index = 0; index < limit && index < mask.length(); index = index + 1 {
    if mask[index] {
      result[index] = original[index]
    }
  }
  result
}

///|
pub fn imputation_reconstruction_error(
  original : Array[Double],
  filled : Array[Double],
  observed_mask : Array[Bool],
) -> Double {
  let errors = []
  let limit = if original.length() < filled.length() {
    original.length()
  } else {
    filled.length()
  }
  for index = 0
      index < limit && index < observed_mask.length()
      index = index + 1 {
    if !observed_mask[index] {
      errors.push(original[index] - filled[index])
    }
  }
  mad(errors)
}

///|
pub fn imputation_smooth(data : Array[Double], window : Int) -> Array[Double] {
  rolling_median(imputation_median(data), if window < 2 { 2 } else { window })
}

///|
pub fn imputation_smooth_difference(
  data : Array[Double],
  window : Int,
) -> Array[Double] {
  feature_difference(imputation_smooth(data, window), 1)
}

///|
pub fn imputation_fill_and_clip(
  data : Array[Double],
  window : Int,
  rule : QualityRule,
) -> Array[Double] {
  quality_clip(
    imputation_with_strategy(
      data,
      imputation_best_strategy(data, window, rule),
      window,
    ),
    rule,
  )
}

///|
pub fn imputation_complete(
  data : Array[Double],
  window : Int,
  rule : QualityRule,
) -> Bool {
  imputation_missing_count(imputation_fill_and_clip(data, window, rule)) == 0
}