///|
fn explain_classification_node(
  node : ClassificationNode,
  features : Array[Double],
  node_id : Int,
  path : Array[DecisionStep],
) -> ClassificationExplanation {
  match node {
    ClassificationLeaf(prediction, probabilities, samples, impurity) => {
      let copied = probabilities.copy()
      {
        prediction,
        probabilities: copied,
        confidence: copied[prediction],
        leaf_id: node_id,
        leaf_samples: samples,
        leaf_impurity: impurity,
        path,
      }
    }
    ClassificationBranch(feature, threshold, left, right, _, _, _) =>
      if features[feature] <= threshold {
        let next_path = path.copy()
        next_path.push(DecisionStep(feature, threshold, GoLeft))
        explain_classification_node(left, features, node_id + 1, next_path)
      } else {
        let next_path = path.copy()
        next_path.push(DecisionStep(feature, threshold, GoRight))
        let right_id = node_id + 1 + classification_node_count(left)
        explain_classification_node(right, features, right_id, next_path)
      }
  }
}

///|
fn explain_regression_node(
  node : RegressionNode,
  features : Array[Double],
  node_id : Int,
  path : Array[DecisionStep],
) -> RegressionExplanation {
  match node {
    RegressionLeaf(prediction, samples, variance) =>
      {
        prediction,
        leaf_id: node_id,
        leaf_samples: samples,
        leaf_variance: variance,
        path,
      }
    RegressionBranch(feature, threshold, left, right, _, _, _) =>
      if features[feature] <= threshold {
        let next_path = path.copy()
        next_path.push(DecisionStep(feature, threshold, GoLeft))
        explain_regression_node(left, features, node_id + 1, next_path)
      } else {
        let next_path = path.copy()
        next_path.push(DecisionStep(feature, threshold, GoRight))
        let right_id = node_id + 1 + regression_node_count(left)
        explain_regression_node(right, features, right_id, next_path)
      }
  }
}

///|
fn collect_classification_rules(
  node : ClassificationNode,
  node_id : Int,
  conditions : Array[RuleCondition],
  rules : Array[ClassificationRule],
) -> Int {
  match node {
    ClassificationLeaf(prediction, probabilities, samples, impurity) => {
      rules.push({
        leaf_id: node_id,
        conditions,
        prediction,
        probabilities: probabilities.copy(),
        samples,
        impurity,
      })
      node_id + 1
    }
    ClassificationBranch(feature, threshold, left, right, _, _, _) => {
      let left_conditions = conditions.copy()
      left_conditions.push(RuleCondition(feature, threshold, GoLeft))
      let next_id = collect_classification_rules(
        left,
        node_id + 1,
        left_conditions,
        rules,
      )
      let right_conditions = conditions.copy()
      right_conditions.push(RuleCondition(feature, threshold, GoRight))
      collect_classification_rules(right, next_id, right_conditions, rules)
    }
  }
}

///|
fn collect_regression_rules(
  node : RegressionNode,
  node_id : Int,
  conditions : Array[RuleCondition],
  rules : Array[RegressionRule],
) -> Int {
  match node {
    RegressionLeaf(prediction, samples, variance) => {
      rules.push({ leaf_id: node_id, conditions, prediction, samples, variance })
      node_id + 1
    }
    RegressionBranch(feature, threshold, left, right, _, _, _) => {
      let left_conditions = conditions.copy()
      left_conditions.push(RuleCondition(feature, threshold, GoLeft))
      let next_id = collect_regression_rules(
        left,
        node_id + 1,
        left_conditions,
        rules,
      )
      let right_conditions = conditions.copy()
      right_conditions.push(RuleCondition(feature, threshold, GoRight))
      collect_regression_rules(right, next_id, right_conditions, rules)
    }
  }
}

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

///|
pub fn ClassificationTree::apply(
  self : ClassificationTree,
  features : Array[Double],
) -> Result[Int, TreeError] {
  match self.explain(features) {
    Ok(value) => Ok(value.leaf_id)
    Err(error) => Err(error)
  }
}

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

///|
pub fn ClassificationTree::rules(
  self : ClassificationTree,
) -> Array[ClassificationRule] {
  let rules : Array[ClassificationRule] = []
  ignore(collect_classification_rules(self.root, 0, [], rules))
  rules
}

///|
pub fn RegressionTree::explain(
  self : RegressionTree,
  features : Array[Double],
) -> Result[RegressionExplanation, TreeError] {
  if features.length() != self.feature_total {
    return Err(
      PredictionFeatureCountMismatch(features.length(), self.feature_total),
    )
  }
  Ok(explain_regression_node(self.root, features, 0, []))
}

///|
pub fn RegressionTree::apply(
  self : RegressionTree,
  features : Array[Double],
) -> Result[Int, TreeError] {
  match self.explain(features) {
    Ok(value) => Ok(value.leaf_id)
    Err(error) => Err(error)
  }
}

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

///|
pub fn RegressionTree::rules(self : RegressionTree) -> Array[RegressionRule] {
  let rules : Array[RegressionRule] = []
  ignore(collect_regression_rules(self.root, 0, [], rules))
  rules
}