///| Breadth-limited candidate tree construction from per-depth draft logits.

///| This is a deterministic policy, not a model: model adapters provide the

///|
/// logits while the policy owns node budgets, widths, and stable tie-breaking.
pub enum PolicyError {
  EmptyDepths
  InvalidWidth
  InvalidBudget
  ProbabilityFailure
} derive(Eq, Debug)

///|
pub struct TreePolicy {
  width : Int
  max_nodes : Int
}

///|
pub fn TreePolicy::new(
  width : Int,
  max_nodes : Int,
) -> Result[TreePolicy, PolicyError] {
  if width <= 0 {
    return Err(InvalidWidth)
  }
  if max_nodes <= 0 {
    return Err(InvalidBudget)
  }
  Ok({ width, max_nodes })
}

///|
fn insert_ranked(
  indices : Array[Int],
  values : Array[Double],
  candidate : Int,
  width : Int,
) -> Unit {
  let mut inserted = false
  for position in 0.. values[indices[position]] ||
        (
          values[candidate] == values[indices[position]] &&
          candidate < indices[position]
        )
      ) {
      indices.insert(position, candidate)
      inserted = true
      break
    }
  }
  if !inserted {
    indices.push(candidate)
  }
  if indices.length() > width {
    ignore(indices.pop())
  }
}

///|
pub fn top_tokens(
  logits : Array[Double],
  width : Int,
) -> Result[Array[(Int, Array[Double])], PolicyError] {
  if width <= 0 {
    return Err(InvalidWidth)
  }
  let distribution = match softmax(logits) {
    Ok(value) => value
    Err(_) => return Err(ProbabilityFailure)
  }
  let ranked : Array[Int] = []
  for index in 0.. Result[DraftTree, PolicyError] {
  if depth_logits.length() == 0 {
    return Err(EmptyDepths)
  }
  let nodes : Array[TreeNode] = []
  let mut parents : Array[Int?] = [None]
  let mut next_id = 0
  for depth_index in 0.. value
      Err(error) => return Err(error)
    }
    let next_parents : Array[Int?] = []
    for parent in parents {
      for candidate in candidates {
        if nodes.length() >= self.max_nodes {
          break
        }
        let (token, distribution) = candidate
        nodes.push({
          id: next_id,
          parent,
          token,
          distribution,
          depth: depth_index + 1,
        })
        next_parents.push(Some(next_id))
        next_id = next_id + 1
      }
      if nodes.length() >= self.max_nodes {
        break
      }
    }
    parents = next_parents
    if parents.length() == 0 {
      break
    }
  }
  match make_tree(prefix, nodes) {
    Ok(value) => Ok(value)
    Err(_) => Err(InvalidBudget)
  }
}