///|
fn validate_classification_config(
  config : ClassificationConfig,
) -> Result[Unit, TreeError] {
  if config.max_depth < 0 {
    return Err(InvalidMaxDepth(config.max_depth))
  }
  if config.min_samples_split < 2 {
    return Err(InvalidMinSamplesSplit(config.min_samples_split))
  }
  if config.min_samples_leaf < 1 {
    return Err(InvalidMinSamplesLeaf(config.min_samples_leaf))
  }
  if !finite_number(config.min_impurity_decrease) ||
    config.min_impurity_decrease < 0.0 {
    return Err(InvalidMinImpurityDecrease(config.min_impurity_decrease))
  }
  Ok(())
}

///|
fn classification_counts(
  dataset : ClassificationDataset,
  indices : Array[Int],
) -> Array[Int] {
  let counts = Array::make(dataset.class_total, 0)
  for index in indices {
    let label = dataset.labels[index]
    counts[label] = counts[label] + 1
  }
  counts
}

///|
fn classification_leaf(
  dataset : ClassificationDataset,
  indices : Array[Int],
  impurity : Double,
) -> ClassificationNode {
  let counts = classification_counts(dataset, indices)
  let mut prediction = 0
  for class_index = 1
      class_index < counts.length()
      class_index = class_index + 1 {
    if counts[class_index] > counts[prediction] {
      prediction = class_index
    }
  }
  let probabilities = counts.map(fn(count) {
    count.to_double() / indices.length().to_double()
  })
  ClassificationLeaf(prediction, probabilities, indices.length(), impurity)
}

///|
fn best_classification_split(
  dataset : ClassificationDataset,
  indices : Array[Int],
  config : ClassificationConfig,
  parent_impurity : Double,
) -> ClassificationSplit? {
  let total = indices.length()
  let total_counts = classification_counts(dataset, indices)
  let mut best : ClassificationSplit? = None
  let mut best_gain = -1.0
  for feature_index = 0
      feature_index < dataset.feature_total
      feature_index = feature_index + 1 {
    let sorted = indices.copy()
    sorted.sort_by(fn(left, right) {
      compare_feature_indices(dataset, feature_index, left, right)
    })
    let left_counts = Array::make(dataset.class_total, 0)
    let right_counts = total_counts.copy()
    for split_index = 1; split_index < total; split_index = split_index + 1 {
      let moved_label = dataset.labels[sorted[split_index - 1]]
      left_counts[moved_label] = left_counts[moved_label] + 1
      right_counts[moved_label] = right_counts[moved_label] - 1
      let left_size = split_index
      let right_size = total - split_index
      if left_size < config.min_samples_leaf ||
        right_size < config.min_samples_leaf {
        continue
      }
      let left_value = dataset.rows[sorted[split_index - 1]][feature_index]
      let right_value = dataset.rows[sorted[split_index]][feature_index]
      if left_value == right_value {
        continue
      }
      let left_impurity = impurity_from_counts(
        left_counts,
        left_size,
        config.criterion,
      )
      let right_impurity = impurity_from_counts(
        right_counts,
        right_size,
        config.criterion,
      )
      let weighted = left_size.to_double() / total.to_double() * left_impurity +
        right_size.to_double() / total.to_double() * right_impurity
      let gain = parent_impurity - weighted
      if gain > best_gain + 0.000000000001 {
        best_gain = gain
        best = Some({
          feature_index,
          threshold: left_value + (right_value - left_value) / 2.0,
          gain,
        })
      }
    }
  }
  best
}

///|
fn build_classification_node(
  dataset : ClassificationDataset,
  indices : Array[Int],
  config : ClassificationConfig,
  depth : Int,
) -> ClassificationNode {
  let counts = classification_counts(dataset, indices)
  let impurity = impurity_from_counts(
    counts,
    indices.length(),
    config.criterion,
  )
  if depth >= config.max_depth ||
    indices.length() < config.min_samples_split ||
    impurity <= 0.000000000001 {
    return classification_leaf(dataset, indices, impurity)
  }
  let split = match
    best_classification_split(dataset, indices, config, impurity) {
    Some(value) => value
    None => return classification_leaf(dataset, indices, impurity)
  }
  if split.gain + 0.000000000001 < config.min_impurity_decrease ||
    split.gain <= 0.000000000001 {
    return classification_leaf(dataset, indices, impurity)
  }
  let left_indices : Array[Int] = []
  let right_indices : Array[Int] = []
  for index in indices {
    if dataset.rows[index][split.feature_index] <= split.threshold {
      left_indices.push(index)
    } else {
      right_indices.push(index)
    }
  }
  if left_indices.is_empty() || right_indices.is_empty() {
    return classification_leaf(dataset, indices, impurity)
  }
  ClassificationBranch(
    split.feature_index,
    split.threshold,
    build_classification_node(dataset, left_indices, config, depth + 1),
    build_classification_node(dataset, right_indices, config, depth + 1),
    indices.length(),
    impurity,
    split.gain,
  )
}

///|
/// Trains one deterministic classification tree.
pub fn train_classifier(
  dataset : ClassificationDataset,
  config : ClassificationConfig,
) -> Result[ClassificationTree, TreeError] {
  match validate_classification_config(config) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  let indices = Array::makei(dataset.row_count(), fn(index) { index })
  Ok({
    root: build_classification_node(dataset, indices, config, 0),
    feature_total: dataset.feature_total,
    class_total: dataset.class_total,
    config,
  })
}

///|
fn classification_prediction(
  node : ClassificationNode,
  features : Array[Double],
) -> Int {
  match node {
    ClassificationLeaf(prediction, _, _, _) => prediction
    ClassificationBranch(feature, threshold, left, right, _, _, _) =>
      if features[feature] <= threshold {
        classification_prediction(left, features)
      } else {
        classification_prediction(right, features)
      }
  }
}

///|
fn classification_probabilities(
  node : ClassificationNode,
  features : Array[Double],
) -> Array[Double] {
  match node {
    ClassificationLeaf(_, probabilities, _, _) => probabilities.copy()
    ClassificationBranch(feature, threshold, left, right, _, _, _) =>
      if features[feature] <= threshold {
        classification_probabilities(left, features)
      } else {
        classification_probabilities(right, features)
      }
  }
}

///|
pub fn ClassificationTree::predict(
  self : ClassificationTree,
  features : Array[Double],
) -> Result[Int, TreeError] {
  if features.length() != self.feature_total {
    return Err(
      PredictionFeatureCountMismatch(features.length(), self.feature_total),
    )
  }
  Ok(classification_prediction(self.root, features))
}

///|
pub fn ClassificationTree::predict_proba(
  self : ClassificationTree,
  features : Array[Double],
) -> Result[Array[Double], TreeError] {
  if features.length() != self.feature_total {
    return Err(
      PredictionFeatureCountMismatch(features.length(), self.feature_total),
    )
  }
  Ok(classification_probabilities(self.root, features))
}

///|
pub fn ClassificationTree::predict_batch(
  self : ClassificationTree,
  rows : Array[Array[Double]],
) -> Result[Array[Int], TreeError] {
  let predictions : Array[Int] = []
  for row in rows {
    match self.predict(row) {
      Ok(value) => predictions.push(value)
      Err(error) => return Err(error)
    }
  }
  Ok(predictions)
}