///|
fn weighted_classification_depth(node : WeightedClassificationNode) -> Int {
  match node {
    WeightedClassificationLeaf(_, _, _, _, _) => 0
    WeightedClassificationBranch(_, _, left, right, _, _, _, _) => {
      let left_depth = weighted_classification_depth(left)
      let right_depth = weighted_classification_depth(right)
      let maximum = if left_depth > right_depth {
        left_depth
      } else {
        right_depth
      }
      maximum + 1
    }
  }
}

///|
fn weighted_classification_node_count(node : WeightedClassificationNode) -> Int {
  match node {
    WeightedClassificationLeaf(_, _, _, _, _) => 1
    WeightedClassificationBranch(_, _, left, right, _, _, _, _) =>
      1 +
      weighted_classification_node_count(left) +
      weighted_classification_node_count(right)
  }
}

///|
fn weighted_classification_leaf_count(node : WeightedClassificationNode) -> Int {
  match node {
    WeightedClassificationLeaf(_, _, _, _, _) => 1
    WeightedClassificationBranch(_, _, left, right, _, _, _, _) =>
      weighted_classification_leaf_count(left) +
      weighted_classification_leaf_count(right)
  }
}

///|
fn weighted_classification_importance(
  node : WeightedClassificationNode,
  values : Array[Double],
) -> Unit {
  match node {
    WeightedClassificationLeaf(_, _, _, _, _) => ()
    WeightedClassificationBranch(feature, _, left, right, _, weight, _, gain) => {
      values[feature] = values[feature] + weight * gain
      weighted_classification_importance(left, values)
      weighted_classification_importance(right, values)
    }
  }
}

///|
fn weighted_classification_path(
  node : WeightedClassificationNode,
  features : Array[Double],
  path : Array[DecisionStep],
) -> Unit {
  match node {
    WeightedClassificationLeaf(_, _, _, _, _) => ()
    WeightedClassificationBranch(feature, threshold, left, right, _, _, _, _) =>
      if features[feature] <= threshold {
        path.push(DecisionStep(feature, threshold, GoLeft))
        weighted_classification_path(left, features, path)
      } else {
        path.push(DecisionStep(feature, threshold, GoRight))
        weighted_classification_path(right, features, path)
      }
  }
}

///|
fn weighted_regression_depth(node : WeightedRegressionNode) -> Int {
  match node {
    WeightedRegressionLeaf(_, _, _, _) => 0
    WeightedRegressionBranch(_, _, left, right, _, _, _, _) => {
      let left_depth = weighted_regression_depth(left)
      let right_depth = weighted_regression_depth(right)
      let maximum = if left_depth > right_depth {
        left_depth
      } else {
        right_depth
      }
      maximum + 1
    }
  }
}

///|
fn weighted_regression_node_count(node : WeightedRegressionNode) -> Int {
  match node {
    WeightedRegressionLeaf(_, _, _, _) => 1
    WeightedRegressionBranch(_, _, left, right, _, _, _, _) =>
      1 +
      weighted_regression_node_count(left) +
      weighted_regression_node_count(right)
  }
}

///|
fn weighted_regression_leaf_count(node : WeightedRegressionNode) -> Int {
  match node {
    WeightedRegressionLeaf(_, _, _, _) => 1
    WeightedRegressionBranch(_, _, left, right, _, _, _, _) =>
      weighted_regression_leaf_count(left) +
      weighted_regression_leaf_count(right)
  }
}

///|
fn weighted_regression_importance(
  node : WeightedRegressionNode,
  values : Array[Double],
) -> Unit {
  match node {
    WeightedRegressionLeaf(_, _, _, _) => ()
    WeightedRegressionBranch(feature, _, left, right, _, weight, _, gain) => {
      values[feature] = values[feature] + weight * gain
      weighted_regression_importance(left, values)
      weighted_regression_importance(right, values)
    }
  }
}

///|
fn weighted_regression_path(
  node : WeightedRegressionNode,
  features : Array[Double],
  path : Array[DecisionStep],
) -> Unit {
  match node {
    WeightedRegressionLeaf(_, _, _, _) => ()
    WeightedRegressionBranch(feature, threshold, left, right, _, _, _, _) =>
      if features[feature] <= threshold {
        path.push(DecisionStep(feature, threshold, GoLeft))
        weighted_regression_path(left, features, path)
      } else {
        path.push(DecisionStep(feature, threshold, GoRight))
        weighted_regression_path(right, features, path)
      }
  }
}

///|
pub fn WeightedClassificationTree::depth(
  self : WeightedClassificationTree,
) -> Int {
  weighted_classification_depth(self.root)
}

///|
pub fn WeightedClassificationTree::node_count(
  self : WeightedClassificationTree,
) -> Int {
  weighted_classification_node_count(self.root)
}

///|
pub fn WeightedClassificationTree::leaf_count(
  self : WeightedClassificationTree,
) -> Int {
  weighted_classification_leaf_count(self.root)
}

///|
pub fn WeightedClassificationTree::feature_importance(
  self : WeightedClassificationTree,
) -> Array[Double] {
  let values = Array::make(self.feature_total, 0.0)
  weighted_classification_importance(self.root, values)
  normalize_importance(values)
}

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

///|
pub fn WeightedRegressionTree::depth(self : WeightedRegressionTree) -> Int {
  weighted_regression_depth(self.root)
}

///|
pub fn WeightedRegressionTree::node_count(self : WeightedRegressionTree) -> Int {
  weighted_regression_node_count(self.root)
}

///|
pub fn WeightedRegressionTree::leaf_count(self : WeightedRegressionTree) -> Int {
  weighted_regression_leaf_count(self.root)
}

///|
pub fn WeightedRegressionTree::feature_importance(
  self : WeightedRegressionTree,
) -> Array[Double] {
  let values = Array::make(self.feature_total, 0.0)
  weighted_regression_importance(self.root, values)
  normalize_importance(values)
}

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