///| Candidate trees generalize a draft sequence: every node extends one parent

///| path, so a target model can score several continuations in one verifier

///| batch. Nodes are stored in parent-before-child order for deterministic

///|
/// traversal and simple validation.
pub enum TreeError {
  EmptyTree
  InvalidRoot
  InvalidParent(Int, Int)
  DuplicateNodeId(Int)
  InvalidNodeDistribution(Int)
  InvalidNodeToken(Int)
  BudgetExceeded(Int, Int)
} derive(Eq, Debug)

///|
pub struct TreeNode {
  id : Int
  parent : Int?
  token : Int
  distribution : Array[Double]
  depth : Int
}

///|
pub struct DraftTree {
  prefix : Array[Int]
  nodes : Array[TreeNode]
}

///|
pub fn DraftTree::node_count(self : DraftTree) -> Int {
  self.nodes.length()
}

///|
pub fn DraftTree::max_depth(self : DraftTree) -> Int {
  let mut maximum = 0
  for node in self.nodes {
    if node.depth > maximum {
      maximum = node.depth
    }
  }
  maximum
}

///|
pub fn DraftTree::node(self : DraftTree, id : Int) -> TreeNode? {
  for node in self.nodes {
    if node.id == id {
      return Some(node)
    }
  }
  None
}

///|
pub fn DraftTree::children(self : DraftTree, parent : Int?) -> Array[TreeNode] {
  let result : Array[TreeNode] = []
  for node in self.nodes {
    match (node.parent, parent) {
      (None, None) => result.push(node)
      (Some(left), Some(right)) if left == right => result.push(node)
      _ => ()
    }
  }
  result
}

///|
pub fn DraftTree::path_to(self : DraftTree, id : Int) -> Array[Int] {
  let reverse : Array[Int] = []
  let mut current = Some(id)
  while current is Some(node_id) {
    match self.node(node_id) {
      None => return []
      Some(node) => {
        reverse.push(node.token)
        current = node.parent
      }
    }
  }
  let path : Array[Int] = []
  for index in 0.. Bool {
  distribution_is_valid(values)
}

///|
pub fn validate_tree(
  tree : DraftTree,
  node_budget? : Int = 1024,
) -> Result[Unit, TreeError] {
  if tree.nodes.length() == 0 {
    return Err(EmptyTree)
  }
  if tree.nodes.length() > node_budget {
    return Err(BudgetExceeded(tree.nodes.length(), node_budget))
  }
  let seen : Array[Int] = []
  for node in tree.nodes {
    for id in seen {
      if id == node.id {
        return Err(DuplicateNodeId(node.id))
      }
    }
    if !valid_node_distribution(node.distribution) {
      return Err(InvalidNodeDistribution(node.id))
    }
    if node.token < 0 || node.token >= node.distribution.length() {
      return Err(InvalidNodeToken(node.id))
    }
    if node.distribution.length() != tree.nodes[0].distribution.length() {
      return Err(InvalidNodeDistribution(node.id))
    }
    for previous in tree.nodes {
      if previous.id != node.id &&
        previous.parent == node.parent &&
        previous.token == node.token {
        return Err(InvalidNodeToken(node.id))
      }
    }
    match node.parent {
      None => if node.depth != 1 { return Err(InvalidRoot) }
      Some(parent_id) => {
        if !contains_id(seen, parent_id) {
          return Err(InvalidParent(node.id, parent_id))
        }
        match tree.node(parent_id) {
          None => return Err(InvalidParent(node.id, parent_id))
          Some(parent) =>
            if parent.depth + 1 != node.depth {
              return Err(InvalidParent(node.id, parent_id))
            }
        }
      }
    }
    seen.push(node.id)
  }
  Ok(())
}

///|
pub fn make_tree(
  prefix : Array[Int],
  nodes : Array[TreeNode],
) -> Result[DraftTree, TreeError] {
  let tree = { prefix, nodes }
  match validate_tree(tree) {
    Ok(_) => Ok(tree)
    Err(error) => Err(error)
  }
}

///|
pub fn prune_tree(
  tree : DraftTree,
  max_depth : Int,
  max_nodes : Int,
) -> DraftTree {
  let nodes : Array[TreeNode] = []
  for node in tree.nodes {
    if node.depth <= max_depth && nodes.length() < max_nodes {
      nodes.push(node)
    }
  }
  { prefix: tree.prefix, nodes }
}