///|
fn validate_weights(
  weights : Array[Double],
  row_count : Int,
) -> Result[(Array[Double], Double), TreeError] {
  if weights.length() != row_count {
    return Err(WeightCountMismatch(row_count, weights.length()))
  }
  let copied = weights.copy()
  let mut total = 0.0
  for index, weight in copied {
    if !finite_number(weight) || weight < 0.0 {
      return Err(InvalidSampleWeight(index))
    }
    total = total + weight
  }
  if total <= 0.0 {
    return Err(ZeroTotalWeight)
  }
  Ok((copied, total))
}

///|
pub fn weighted_classification_dataset(
  rows : Array[Array[Double]],
  labels : Array[Int],
  weights : Array[Double],
  class_count : Int,
) -> Result[WeightedClassificationDataset, TreeError] {
  if class_count <= 0 {
    return Err(InvalidClassCount(class_count))
  }
  let (copied_rows, feature_count) = match validate_rows(rows) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  if labels.length() != copied_rows.length() {
    return Err(TargetCountMismatch(copied_rows.length(), labels.length()))
  }
  for index, label in labels {
    if label < 0 || label >= class_count {
      return Err(InvalidClassLabel(index, label, class_count))
    }
  }
  let (copied_weights, weight_total) = match
    validate_weights(weights, copied_rows.length()) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  Ok({
    rows: copied_rows,
    labels: labels.copy(),
    weights: copied_weights,
    feature_total: feature_count,
    class_total: class_count,
    weight_total,
  })
}

///|
pub fn weighted_regression_dataset(
  rows : Array[Array[Double]],
  targets : Array[Double],
  weights : Array[Double],
) -> Result[WeightedRegressionDataset, TreeError] {
  let (copied_rows, feature_count) = match validate_rows(rows) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  if targets.length() != copied_rows.length() {
    return Err(TargetCountMismatch(copied_rows.length(), targets.length()))
  }
  for index, target in targets {
    if !finite_number(target) {
      return Err(NonFiniteTarget(index))
    }
  }
  let (copied_weights, weight_total) = match
    validate_weights(weights, copied_rows.length()) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  Ok({
    rows: copied_rows,
    targets: targets.copy(),
    weights: copied_weights,
    feature_total: feature_count,
    weight_total,
  })
}

///|
pub fn WeightedClassificationDataset::row_count(
  self : WeightedClassificationDataset,
) -> Int {
  self.rows.length()
}

///|
pub fn WeightedClassificationDataset::feature_count(
  self : WeightedClassificationDataset,
) -> Int {
  self.feature_total
}

///|
pub fn WeightedClassificationDataset::class_count(
  self : WeightedClassificationDataset,
) -> Int {
  self.class_total
}

///|
pub fn WeightedClassificationDataset::total_weight(
  self : WeightedClassificationDataset,
) -> Double {
  self.weight_total
}

///|
pub fn WeightedRegressionDataset::row_count(
  self : WeightedRegressionDataset,
) -> Int {
  self.rows.length()
}

///|
pub fn WeightedRegressionDataset::feature_count(
  self : WeightedRegressionDataset,
) -> Int {
  self.feature_total
}

///|
pub fn WeightedRegressionDataset::total_weight(
  self : WeightedRegressionDataset,
) -> Double {
  self.weight_total
}

///|
fn weighted_impurity(
  counts : Array[Double],
  weight_total : Double,
  criterion : ClassificationCriterion,
) -> Double {
  if weight_total <= 0.0 {
    return 0.0
  }
  match criterion {
    Gini => {
      let mut result = 1.0
      for count in counts {
        if count > 0.0 {
          let probability = count / weight_total
          result = result - probability * probability
        }
      }
      result
    }
    Entropy => {
      let mut result = 0.0
      for count in counts {
        if count > 0.0 {
          let probability = count / weight_total
          result = result - probability * log2_positive(probability)
        }
      }
      result
    }
  }
}

///|
fn weighted_class_counts(
  dataset : WeightedClassificationDataset,
  indices : Array[Int],
) -> (Array[Double], Double) {
  let counts = Array::make(dataset.class_total, 0.0)
  let mut weight_total = 0.0
  for index in indices {
    let weight = dataset.weights[index]
    counts[dataset.labels[index]] = counts[dataset.labels[index]] + weight
    weight_total = weight_total + weight
  }
  (counts, weight_total)
}

///|
fn weighted_classification_leaf(
  dataset : WeightedClassificationDataset,
  indices : Array[Int],
  impurity : Double,
) -> WeightedClassificationNode {
  let (counts, weight_total) = weighted_class_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 / weight_total })
  WeightedClassificationLeaf(
    prediction,
    probabilities,
    indices.length(),
    weight_total,
    impurity,
  )
}

///|
fn compare_weighted_classification_indices(
  dataset : WeightedClassificationDataset,
  feature_index : Int,
  left : Int,
  right : Int,
) -> Int {
  let left_value = dataset.rows[left][feature_index]
  let right_value = dataset.rows[right][feature_index]
  if left_value < right_value {
    -1
  } else if left_value > right_value {
    1
  } else if left < right {
    -1
  } else if left > right {
    1
  } else {
    0
  }
}

///|
fn best_weighted_classification_split(
  dataset : WeightedClassificationDataset,
  indices : Array[Int],
  config : ClassificationConfig,
  parent_impurity : Double,
  total_weight : Double,
) -> WeightedClassificationSplit? {
  let total = indices.length()
  let (total_counts, _) = weighted_class_counts(dataset, indices)
  let mut best : WeightedClassificationSplit? = 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_weighted_classification_indices(
        dataset, feature_index, left, right,
      )
    })
    let left_counts = Array::make(dataset.class_total, 0.0)
    let right_counts = total_counts.copy()
    let mut left_weight = 0.0
    let mut right_weight = total_weight
    for split_index = 1; split_index < total; split_index = split_index + 1 {
      let moved = sorted[split_index - 1]
      let moved_label = dataset.labels[moved]
      let moved_weight = dataset.weights[moved]
      left_counts[moved_label] = left_counts[moved_label] + moved_weight
      right_counts[moved_label] = right_counts[moved_label] - moved_weight
      left_weight = left_weight + moved_weight
      right_weight = right_weight - moved_weight
      let left_size = split_index
      let right_size = total - split_index
      if left_size < config.min_samples_leaf ||
        right_size < config.min_samples_leaf ||
        left_weight <= 0.0 ||
        right_weight <= 0.0 {
        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 = weighted_impurity(
        left_counts,
        left_weight,
        config.criterion,
      )
      let right_impurity = weighted_impurity(
        right_counts,
        right_weight,
        config.criterion,
      )
      let weighted = left_weight / total_weight * left_impurity +
        right_weight / total_weight * 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_weighted_classification_node(
  dataset : WeightedClassificationDataset,
  indices : Array[Int],
  config : ClassificationConfig,
  depth : Int,
) -> WeightedClassificationNode {
  let (counts, weight_total) = weighted_class_counts(dataset, indices)
  let impurity = weighted_impurity(counts, weight_total, config.criterion)
  if depth >= config.max_depth ||
    indices.length() < config.min_samples_split ||
    impurity <= 0.000000000001 {
    return weighted_classification_leaf(dataset, indices, impurity)
  }
  let split = match
    best_weighted_classification_split(
      dataset, indices, config, impurity, weight_total,
    ) {
    Some(value) => value
    None => return weighted_classification_leaf(dataset, indices, impurity)
  }
  if split.gain + 0.000000000001 < config.min_impurity_decrease ||
    split.gain <= 0.000000000001 {
    return weighted_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)
    }
  }
  WeightedClassificationBranch(
    split.feature_index,
    split.threshold,
    build_weighted_classification_node(dataset, left_indices, config, depth + 1),
    build_weighted_classification_node(
      dataset,
      right_indices,
      config,
      depth + 1,
    ),
    indices.length(),
    weight_total,
    impurity,
    split.gain,
  )
}

///|
pub fn train_weighted_classifier(
  dataset : WeightedClassificationDataset,
  config : ClassificationConfig,
) -> Result[WeightedClassificationTree, 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_weighted_classification_node(dataset, indices, config, 0),
    feature_total: dataset.feature_total,
    class_total: dataset.class_total,
    config,
  })
}

///|
fn weighted_classification_prediction(
  node : WeightedClassificationNode,
  features : Array[Double],
) -> Int {
  match node {
    WeightedClassificationLeaf(prediction, _, _, _, _) => prediction
    WeightedClassificationBranch(feature, threshold, left, right, _, _, _, _) =>
      if features[feature] <= threshold {
        weighted_classification_prediction(left, features)
      } else {
        weighted_classification_prediction(right, features)
      }
  }
}

///|
fn weighted_classification_probabilities(
  node : WeightedClassificationNode,
  features : Array[Double],
) -> Array[Double] {
  match node {
    WeightedClassificationLeaf(_, probabilities, _, _, _) =>
      probabilities.copy()
    WeightedClassificationBranch(feature, threshold, left, right, _, _, _, _) =>
      if features[feature] <= threshold {
        weighted_classification_probabilities(left, features)
      } else {
        weighted_classification_probabilities(right, features)
      }
  }
}

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

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

///|
pub fn WeightedClassificationTree::predict_batch(
  self : WeightedClassificationTree,
  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)
}

///|
fn weighted_regression_stats(
  dataset : WeightedRegressionDataset,
  indices : Array[Int],
) -> (Double, Double, Double) {
  let mut weight_total = 0.0
  let mut weighted_sum = 0.0
  let mut weighted_squared_sum = 0.0
  for index in indices {
    let weight = dataset.weights[index]
    let target = dataset.targets[index]
    weight_total = weight_total + weight
    weighted_sum = weighted_sum + weight * target
    weighted_squared_sum = weighted_squared_sum + weight * target * target
  }
  (weight_total, weighted_sum, weighted_squared_sum)
}

///|
fn weighted_variance(
  weight_total : Double,
  weighted_sum : Double,
  weighted_squared_sum : Double,
) -> Double {
  if weight_total <= 0.0 {
    return 0.0
  }
  let mean = weighted_sum / weight_total
  let value = weighted_squared_sum / weight_total - mean * mean
  if value > 0.0 {
    value
  } else {
    0.0
  }
}

///|
fn weighted_regression_leaf(
  dataset : WeightedRegressionDataset,
  indices : Array[Int],
) -> WeightedRegressionNode {
  let (weight_total, weighted_sum, weighted_squared_sum) = weighted_regression_stats(
    dataset, indices,
  )
  WeightedRegressionLeaf(
    weighted_sum / weight_total,
    indices.length(),
    weight_total,
    weighted_variance(weight_total, weighted_sum, weighted_squared_sum),
  )
}

///|
fn compare_weighted_regression_indices(
  dataset : WeightedRegressionDataset,
  feature_index : Int,
  left : Int,
  right : Int,
) -> Int {
  let left_value = dataset.rows[left][feature_index]
  let right_value = dataset.rows[right][feature_index]
  if left_value < right_value {
    -1
  } else if left_value > right_value {
    1
  } else if left < right {
    -1
  } else if left > right {
    1
  } else {
    0
  }
}

///|
fn best_weighted_regression_split(
  dataset : WeightedRegressionDataset,
  indices : Array[Int],
  config : RegressionConfig,
  parent_variance : Double,
  total_weight : Double,
) -> WeightedRegressionSplit? {
  let total = indices.length()
  let (_, total_sum, total_squared_sum) = weighted_regression_stats(
    dataset, indices,
  )
  let mut best : WeightedRegressionSplit? = 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_weighted_regression_indices(dataset, feature_index, left, right)
    })
    let mut left_weight = 0.0
    let mut left_sum = 0.0
    let mut left_squared_sum = 0.0
    let mut right_weight = total_weight
    let mut right_sum = total_sum
    let mut right_squared_sum = total_squared_sum
    for split_index = 1; split_index < total; split_index = split_index + 1 {
      let moved = sorted[split_index - 1]
      let weight = dataset.weights[moved]
      let target = dataset.targets[moved]
      left_weight = left_weight + weight
      left_sum = left_sum + weight * target
      left_squared_sum = left_squared_sum + weight * target * target
      right_weight = right_weight - weight
      right_sum = right_sum - weight * target
      right_squared_sum = right_squared_sum - weight * target * target
      let left_size = split_index
      let right_size = total - split_index
      if left_size < config.min_samples_leaf ||
        right_size < config.min_samples_leaf ||
        left_weight <= 0.0 ||
        right_weight <= 0.0 {
        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_variance = weighted_variance(
        left_weight, left_sum, left_squared_sum,
      )
      let right_variance = weighted_variance(
        right_weight, right_sum, right_squared_sum,
      )
      let combined = left_weight / total_weight * left_variance +
        right_weight / total_weight * right_variance
      let gain = parent_variance - combined
      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_weighted_regression_node(
  dataset : WeightedRegressionDataset,
  indices : Array[Int],
  config : RegressionConfig,
  depth : Int,
) -> WeightedRegressionNode {
  let (weight_total, weighted_sum, weighted_squared_sum) = weighted_regression_stats(
    dataset, indices,
  )
  let variance = weighted_variance(
    weight_total, weighted_sum, weighted_squared_sum,
  )
  if depth >= config.max_depth ||
    indices.length() < config.min_samples_split ||
    variance <= 0.000000000001 {
    return weighted_regression_leaf(dataset, indices)
  }
  let split = match
    best_weighted_regression_split(
      dataset, indices, config, variance, weight_total,
    ) {
    Some(value) => value
    None => return weighted_regression_leaf(dataset, indices)
  }
  if split.gain + 0.000000000001 < config.min_impurity_decrease ||
    split.gain <= 0.000000000001 {
    return weighted_regression_leaf(dataset, indices)
  }
  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)
    }
  }
  WeightedRegressionBranch(
    split.feature_index,
    split.threshold,
    build_weighted_regression_node(dataset, left_indices, config, depth + 1),
    build_weighted_regression_node(dataset, right_indices, config, depth + 1),
    indices.length(),
    weight_total,
    variance,
    split.gain,
  )
}

///|
pub fn train_weighted_regressor(
  dataset : WeightedRegressionDataset,
  config : RegressionConfig,
) -> Result[WeightedRegressionTree, TreeError] {
  match validate_regression_config(config) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  let indices = Array::makei(dataset.row_count(), fn(index) { index })
  Ok({
    root: build_weighted_regression_node(dataset, indices, config, 0),
    feature_total: dataset.feature_total,
    config,
  })
}

///|
fn weighted_regression_prediction(
  node : WeightedRegressionNode,
  features : Array[Double],
) -> Double {
  match node {
    WeightedRegressionLeaf(prediction, _, _, _) => prediction
    WeightedRegressionBranch(feature, threshold, left, right, _, _, _, _) =>
      if features[feature] <= threshold {
        weighted_regression_prediction(left, features)
      } else {
        weighted_regression_prediction(right, features)
      }
  }
}

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

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

///|
fn validate_metric_weights(
  weights : Array[Double],
  count : Int,
) -> Result[Double, TreeError] {
  match validate_weights(weights, count) {
    Ok((_, total)) => Ok(total)
    Err(error) => Err(error)
  }
}

///|
pub fn weighted_classification_accuracy(
  actual : Array[Int],
  predicted : Array[Int],
  weights : Array[Double],
) -> Result[Double, TreeError] {
  match validate_metric_lengths(actual.length(), predicted.length()) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  let total_weight = match validate_metric_weights(weights, actual.length()) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  let mut correct_weight = 0.0
  for index = 0; index < actual.length(); index = index + 1 {
    if actual[index] == predicted[index] {
      correct_weight = correct_weight + weights[index]
    }
  }
  Ok(correct_weight / total_weight)
}

///|
pub fn weighted_mean_squared_error(
  actual : Array[Double],
  predicted : Array[Double],
  weights : Array[Double],
) -> Result[Double, TreeError] {
  match validate_metric_lengths(actual.length(), predicted.length()) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  let total_weight = match validate_metric_weights(weights, actual.length()) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  let mut total = 0.0
  for index = 0; index < actual.length(); index = index + 1 {
    let difference = actual[index] - predicted[index]
    total = total + weights[index] * difference * difference
  }
  Ok(total / total_weight)
}

///|
pub fn weighted_mean_absolute_error(
  actual : Array[Double],
  predicted : Array[Double],
  weights : Array[Double],
) -> Result[Double, TreeError] {
  match validate_metric_lengths(actual.length(), predicted.length()) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  let total_weight = match validate_metric_weights(weights, actual.length()) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  let mut total = 0.0
  for index = 0; index < actual.length(); index = index + 1 {
    let difference = actual[index] - predicted[index]
    let absolute = if difference < 0.0 { -difference } else { difference }
    total = total + weights[index] * absolute
  }
  Ok(total / total_weight)
}