///|
fn classification_depth(node : ClassificationNode) -> Int {
  match node {
    ClassificationLeaf(_, _, _, _) => 0
    ClassificationBranch(_, _, left, right, _, _, _) => {
      let left_depth = classification_depth(left)
      let right_depth = classification_depth(right)
      let maximum = if left_depth > right_depth {
        left_depth
      } else {
        right_depth
      }
      1 + maximum
    }
  }
}

///|
fn classification_node_count(node : ClassificationNode) -> Int {
  match node {
    ClassificationLeaf(_, _, _, _) => 1
    ClassificationBranch(_, _, left, right, _, _, _) =>
      1 + classification_node_count(left) + classification_node_count(right)
  }
}

///|
fn classification_leaf_count(node : ClassificationNode) -> Int {
  match node {
    ClassificationLeaf(_, _, _, _) => 1
    ClassificationBranch(_, _, left, right, _, _, _) =>
      classification_leaf_count(left) + classification_leaf_count(right)
  }
}

///|
fn classification_importance(
  node : ClassificationNode,
  values : Array[Double],
) -> Unit {
  match node {
    ClassificationLeaf(_, _, _, _) => ()
    ClassificationBranch(feature, _, left, right, samples, _, gain) => {
      values[feature] = values[feature] + samples.to_double() * gain
      classification_importance(left, values)
      classification_importance(right, values)
    }
  }
}

///|
fn regression_depth(node : RegressionNode) -> Int {
  match node {
    RegressionLeaf(_, _, _) => 0
    RegressionBranch(_, _, left, right, _, _, _) => {
      let left_depth = regression_depth(left)
      let right_depth = regression_depth(right)
      let maximum = if left_depth > right_depth {
        left_depth
      } else {
        right_depth
      }
      1 + maximum
    }
  }
}

///|
fn regression_node_count(node : RegressionNode) -> Int {
  match node {
    RegressionLeaf(_, _, _) => 1
    RegressionBranch(_, _, left, right, _, _, _) =>
      1 + regression_node_count(left) + regression_node_count(right)
  }
}

///|
fn regression_leaf_count(node : RegressionNode) -> Int {
  match node {
    RegressionLeaf(_, _, _) => 1
    RegressionBranch(_, _, left, right, _, _, _) =>
      regression_leaf_count(left) + regression_leaf_count(right)
  }
}

///|
fn regression_importance(node : RegressionNode, values : Array[Double]) -> Unit {
  match node {
    RegressionLeaf(_, _, _) => ()
    RegressionBranch(feature, _, left, right, samples, _, gain) => {
      values[feature] = values[feature] + samples.to_double() * gain
      regression_importance(left, values)
      regression_importance(right, values)
    }
  }
}

///|
fn normalize_importance(values : Array[Double]) -> Array[Double] {
  let mut total = 0.0
  for value in values {
    total = total + value
  }
  if total > 0.0 {
    for index = 0; index < values.length(); index = index + 1 {
      values[index] = values[index] / total
    }
  }
  values
}

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

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

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

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

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

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

///|
pub fn ClassificationTree::decision_path(
  self : ClassificationTree,
  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] = []
  classification_path(self.root, features, path)
  Ok(path)
}

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

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

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

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

///|
pub fn RegressionTree::decision_path(
  self : RegressionTree,
  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] = []
  regression_path(self.root, features, path)
  Ok(path)
}