///| Entropy-adaptive tree construction. Fixed-width trees spend the same

///| branching budget on confident and uncertain draft distributions; this

///| policy expands ambiguous levels and narrows confident ones while retaining

///|
/// an explicit global node budget.
pub enum AdaptiveTreeError {
  InvalidWidth
  InvalidDepth
  InvalidBudget
  InvalidEntropyThreshold
  EmptyLogitSchedule
  ProbabilityFailure
} derive(Eq, Debug)

///| Mutable depth responds to observed draft acceptance. Width is chosen for

///| each logit row from normalized entropy, so the two controls address

///|
/// different properties: expected draft quality and local token ambiguity.
pub struct AdaptiveTreePolicy {
  min_width : Int
  max_width : Int
  min_depth : Int
  max_depth : Int
  max_nodes : Int
  entropy_threshold : Double
  target_acceptance : Double
  mut depth : Int
}

///|
pub fn AdaptiveTreePolicy::new(
  min_width : Int,
  max_width : Int,
  min_depth : Int,
  max_depth : Int,
  max_nodes : Int,
  entropy_threshold : Double,
  target_acceptance : Double,
) -> Result[AdaptiveTreePolicy, AdaptiveTreeError] {
  if min_width <= 0 || max_width < min_width {
    return Err(InvalidWidth)
  }
  if min_depth <= 0 || max_depth < min_depth {
    return Err(InvalidDepth)
  }
  if max_nodes <= 0 {
    return Err(InvalidBudget)
  }
  if entropy_threshold < 0.0 ||
    entropy_threshold > 1.0 ||
    target_acceptance < 0.0 ||
    target_acceptance > 1.0 {
    return Err(InvalidEntropyThreshold)
  }
  Ok({
    min_width,
    max_width,
    min_depth,
    max_depth,
    max_nodes,
    entropy_threshold,
    target_acceptance,
    depth: min_depth,
  })
}

///|
pub fn AdaptiveTreePolicy::depth(self : AdaptiveTreePolicy) -> Int {
  self.depth
}

///| Shannon entropy divided by the maximum entropy for this vocabulary. A

///| one-token vocabulary is treated as perfectly confident to avoid division

///|
/// by log(1).
pub fn normalized_entropy(
  logits : Array[Double],
) -> Result[Double, AdaptiveTreeError] {
  let distribution = match softmax(logits) {
    Ok(value) => value
    Err(_) => return Err(ProbabilityFailure)
  }
  if distribution.length() <= 1 {
    return Ok(0.0)
  }
  let mut entropy = 0.0
  for probability in distribution {
    if probability > 0.0 {
      entropy = entropy - probability * @math.ln(probability)
    }
  }
  Ok(entropy / @math.ln(distribution.length().to_double()))
}

///| Choose a row's branching width. The interpolation is deliberately

///| monotonic and integer-rounded so small entropy changes do not produce an

///|
/// unbounded number of policy states.
pub fn AdaptiveTreePolicy::width_for_logits(
  self : AdaptiveTreePolicy,
  logits : Array[Double],
) -> Result[Int, AdaptiveTreeError] {
  let entropy = match normalized_entropy(logits) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  if entropy <= self.entropy_threshold || self.max_width == self.min_width {
    return Ok(self.min_width)
  }
  let span = self.max_width - self.min_width
  let denominator = 1.0 - self.entropy_threshold
  let fraction = ((entropy - self.entropy_threshold) / denominator).min(1.0)
  Ok(self.min_width + (fraction * span.to_double()).to_int())
}

///| Update planned proposal depth after an observed verifier outcome. This

///|
/// does not change width: width remains tied to next-step uncertainty.
pub fn AdaptiveTreePolicy::observe(
  self : AdaptiveTreePolicy,
  accepted : Int,
  proposed : Int,
) -> Unit {
  if proposed <= 0 {
    return
  }
  let rate = accepted.to_double() / proposed.to_double()
  if rate >= self.target_acceptance && self.depth < self.max_depth {
    self.depth = self.depth + 1
  } else if rate < self.target_acceptance && self.depth > self.min_depth {
    self.depth = self.depth - 1
  }
}

///| Build an entropy-adaptive tree from at most the currently planned depth.

///| Each depth row shares its candidate ranking across parents, matching the

///| simple logit-only reference model. Real adapters may supply context-aware

///|
/// rows per parent, then use `plan_tree_batch` to execute them.
pub fn AdaptiveTreePolicy::build(
  self : AdaptiveTreePolicy,
  prefix : Array[Int],
  depth_logits : Array[Array[Double]],
) -> Result[DraftTree, AdaptiveTreeError] {
  if depth_logits.length() == 0 {
    return Err(EmptyLogitSchedule)
  }
  let nodes : Array[TreeNode] = []
  let mut parents : Array[Int?] = [None]
  let mut next_id = 0
  let usable_depth = self.depth.min(depth_logits.length())
  for depth_index in 0.. value
      Err(error) => return Err(error)
    }
    let candidates = match top_tokens(depth_logits[depth_index], width) {
      Ok(value) => value
      Err(_) => return Err(ProbabilityFailure)
    }
    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(tree) => Ok(tree)
    Err(_) => Err(InvalidBudget)
  }
}

///|
/// Explain the active policy state in a stable format for benchmark metadata.
pub fn AdaptiveTreePolicy::describe(self : AdaptiveTreePolicy) -> String {
  "depth=" +
  self.depth.to_string() +
  " width=" +
  self.min_width.to_string() +
  ".." +
  self.max_width.to_string() +
  " nodes=" +
  self.max_nodes.to_string() +
  " entropy_threshold=" +
  self.entropy_threshold.to_string()
}