///|
/// Both sides of a comparison use the same pure target provider, prompt, and
/// output budget. Callbacks are evaluated sequentially; measured query counts
/// are separate from hypothetical parallel target batches. No latency claim
/// can be inferred from a reduction in batches alone.
pub enum TreeExperimentError {
  InvalidConfiguration
  ModelFailure(ModelTreeError)
  TargetFailure(String)
  InvalidProbabilities
  VerificationFailure(TreeEvaluateError)
  InvalidRandomSeed
} derive(Eq, Debug)

///|
pub struct TreeExperimentConfig {
  width : Int
  depth : Int
  node_budget : Int
  output_tokens : Int
  seed : Int
}

///|
pub fn TreeExperimentConfig::new(
  width : Int,
  depth : Int,
  node_budget : Int,
  output_tokens : Int,
  seed : Int,
) -> Result[TreeExperimentConfig, TreeExperimentError] {
  if width <= 0 ||
    depth <= 0 ||
    depth > 1024 ||
    node_budget <= 0 ||
    node_budget > 1024 ||
    output_tokens <= 0 ||
    output_tokens > 100000 {
    return Err(InvalidConfiguration)
  }
  match DeterministicRng::new(seed) {
    Err(_) => return Err(InvalidRandomSeed)
    Ok(_) => ()
  }
  Ok({ width, depth, node_budget, output_tokens, seed })
}

///|
/// Generated contains only new tokens, excluding the prompt.
pub struct DecodingRun {
  generated : Array[Int]
  target_queries : Int
  draft_queries : Int
  logical_target_batches : Int
  candidate_nodes : Int
  accepted_nodes : Int
}

///|
/// A batch-decoding run together with the proposal depth selected before each
/// round. The trace makes adaptive behavior inspectable without exposing a
/// mutable policy to callers after decoding has finished.
pub struct AdaptiveDecodingRun {
  run : DecodingRun
  planned_depths : Array[Int]
}

///|
pub fn AdaptiveDecodingRun::output_tokens(self : AdaptiveDecodingRun) -> Int {
  self.run.generated.length()
}

///|
pub fn AdaptiveDecodingRun::target_requests(self : AdaptiveDecodingRun) -> Int {
  self.run.target_queries
}

///|
pub fn AdaptiveDecodingRun::planned_depths(
  self : AdaptiveDecodingRun,
) -> Array[Int] {
  self.planned_depths.copy()
}

///|
pub struct TreeExperiment {
  prompt : Array[Int]
  config : TreeExperimentConfig
  baseline : DecodingRun
  speculative : DecodingRun
}

///|
/// The model callback must be a deterministic function of its input tokens.
/// Randomness for token sampling is owned by the caller, not by the model.
pub fn decode_tree_from_model(
  prefix : Array[Int],
  draft : (Array[Int]) -> Result[Array[Double], String],
  target : (Array[Int]) -> Result[Array[Double], String],
  config : TreeExperimentConfig,
) -> Result[DecodingRun, TreeExperimentError] {
  let rng = match DeterministicRng::new(config.seed) {
    Ok(value) => value
    Err(_) => return Err(InvalidRandomSeed)
  }
  let policy = match TreePolicy::new(config.width, config.node_budget) {
    Ok(value) => value
    Err(_) => return Err(InvalidConfiguration)
  }
  let history = prefix.copy()
  let generated : Array[Int] = []
  let mut target_queries = 0
  let mut draft_queries = 0
  let mut batches = 0
  let mut candidate_nodes = 0
  let mut accepted_nodes = 0
  while generated.length() < config.output_tokens {
    let remaining = config.output_tokens - generated.length()
    let proposed = match
      policy.build_from_model(history, config.depth.min(remaining), draft) {
      Ok(value) => value
      Err(error) => return Err(ModelFailure(error))
    }
    let scores = match score_tree_from_model(proposed.tree, target) {
      Ok(value) => value
      Err(error) => return Err(ModelFailure(error))
    }
    let uniforms : Array[Double] = []
    for _ in proposed.tree.nodes {
      uniforms.push(rng.next_unit())
    }
    let fallbacks : Array[Double] = []
    for _ in 0.. value
      Err(error) => return Err(VerificationFailure(error))
    }
    for token in result.emitted {
      history.push(token)
      generated.push(token)
    }
    batches = batches + 1
    target_queries = target_queries + scores.provider_calls
    draft_queries = draft_queries + proposed.provider_calls
    candidate_nodes = candidate_nodes + proposed.tree.node_count()
    accepted_nodes = accepted_nodes + result.accepted_nodes.length()
  }
  Ok({
    generated,
    target_queries,
    draft_queries,
    logical_target_batches: batches,
    candidate_nodes,
    accepted_nodes,
  })
}

///|
/// End-to-end tree decoding through batch providers. Draft expansion submits
/// one frontier batch per depth; target scoring submits one batch per round.
/// `target_queries` and `draft_queries` count provider invocations, not token
/// rows, so callers can compare transport work with the sequential adapter.
pub fn decode_tree_from_batch_models(
  prefix : Array[Int],
  draft : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
  target : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
  config : TreeExperimentConfig,
) -> Result[DecodingRun, TreeExperimentError] {
  let rng = match DeterministicRng::new(config.seed) {
    Ok(value) => value
    Err(_) => return Err(InvalidRandomSeed)
  }
  let policy = match TreePolicy::new(config.width, config.node_budget) {
    Ok(value) => value
    Err(_) => return Err(InvalidConfiguration)
  }
  let history = prefix.copy()
  let generated : Array[Int] = []
  let mut target_queries = 0
  let mut draft_queries = 0
  let mut batches = 0
  let mut candidate_nodes = 0
  let mut accepted_nodes = 0
  while generated.length() < config.output_tokens {
    let remaining = config.output_tokens - generated.length()
    let proposed = match
      policy.build_from_batch_model(history, config.depth.min(remaining), draft) {
      Ok(value) => value
      Err(error) => return Err(ModelFailure(error))
    }
    let scores = match score_tree_from_batch_model(proposed.tree, target) {
      Ok(value) => value
      Err(error) => return Err(ModelFailure(error))
    }
    let uniforms : Array[Double] = []
    for _ in proposed.tree.nodes {
      uniforms.push(rng.next_unit())
    }
    let fallbacks : Array[Double] = []
    for _ in 0.. value
      Err(error) => return Err(VerificationFailure(error))
    }
    for token in result.emitted {
      history.push(token)
      generated.push(token)
    }
    batches = batches + 1
    target_queries = target_queries + scores.provider_calls
    draft_queries = draft_queries + proposed.provider_calls
    candidate_nodes = candidate_nodes + proposed.tree.node_count()
    accepted_nodes = accepted_nodes + result.accepted_nodes.length()
  }
  Ok({
    generated,
    target_queries,
    draft_queries,
    logical_target_batches: batches,
    candidate_nodes,
    accepted_nodes,
  })
}

///|
/// Decode through batch callbacks using the entropy-adaptive tree policy.
/// Width is selected independently for each draft parent row. After a round,
/// the next depth moves toward the observed accepted-path fraction; rejected
/// siblings are deliberately not counted as accepted tokens.
pub fn decode_adaptive_tree_from_batch_models(
  prefix : Array[Int],
  draft : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
  target : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
  policy : AdaptiveTreePolicy,
  output_tokens : Int,
  seed : Int,
) -> Result[AdaptiveDecodingRun, TreeExperimentError] {
  if output_tokens <= 0 || output_tokens > 100000 {
    return Err(InvalidConfiguration)
  }
  let rng = match DeterministicRng::new(seed) {
    Ok(value) => value
    Err(_) => return Err(InvalidRandomSeed)
  }
  let history = prefix.copy()
  let generated : Array[Int] = []
  let depths : Array[Int] = []
  let mut target_queries = 0
  let mut draft_queries = 0
  let mut batches = 0
  let mut candidate_nodes = 0
  let mut accepted_nodes = 0
  while generated.length() < output_tokens {
    let remaining = output_tokens - generated.length()
    let before_depth = policy.depth()
    depths.push(before_depth)
    let proposed = match
      policy.build_from_batch_model_with_depth_limit(history, remaining, draft) {
      Ok(value) => value
      Err(error) => return Err(ModelFailure(error))
    }
    let scores = match score_tree_from_batch_model(proposed.tree, target) {
      Ok(value) => value
      Err(error) => return Err(ModelFailure(error))
    }
    let uniforms : Array[Double] = []
    for _ in proposed.tree.nodes {
      uniforms.push(rng.next_unit())
    }
    let fallbacks : Array[Double] = []
    for _ in 0.. value
      Err(error) => return Err(VerificationFailure(error))
    }
    for token in evaluation.emitted {
      if generated.length() < output_tokens {
        history.push(token)
        generated.push(token)
      }
    }
    policy.observe(
      evaluation.accepted_nodes.length(),
      proposed.tree.max_depth(),
    )
    batches = batches + 1
    target_queries = target_queries + scores.provider_calls
    draft_queries = draft_queries + proposed.provider_calls
    candidate_nodes = candidate_nodes + proposed.tree.node_count()
    accepted_nodes = accepted_nodes + evaluation.accepted_nodes.length()
  }
  Ok({
    run: {
      generated,
      target_queries,
      draft_queries,
      logical_target_batches: batches,
      candidate_nodes,
      accepted_nodes,
    },
    planned_depths: depths,
  })
}

///|
/// Run an ordinary autoregressive baseline against the same model interface.
pub fn decode_baseline_from_model(
  prefix : Array[Int],
  target : (Array[Int]) -> Result[Array[Double], String],
  config : TreeExperimentConfig,
) -> Result[DecodingRun, TreeExperimentError] {
  let rng = match DeterministicRng::new(config.seed) {
    Ok(value) => value
    Err(_) => return Err(InvalidRandomSeed)
  }
  let history = prefix.copy()
  let generated : Array[Int] = []
  let mut vocabulary = 0
  for _ in 0.. value
      Err(error) => return Err(TargetFailure(error))
    }
    if vocabulary == 0 {
      vocabulary = logits.length()
    }
    if vocabulary != logits.length() {
      return Err(InvalidProbabilities)
    }
    let distribution = match softmax(logits) {
      Ok(value) => value
      Err(_) => return Err(InvalidProbabilities)
    }
    let token = match sample_categorical(distribution, rng.next_unit()) {
      Ok(value) => value
      Err(_) => return Err(InvalidProbabilities)
    }
    history.push(token)
    generated.push(token)
  }
  Ok({
    generated,
    target_queries: config.output_tokens,
    draft_queries: 0,
    logical_target_batches: config.output_tokens,
    candidate_nodes: 0,
    accepted_nodes: 0,
  })
}

///|
/// A single entry point fixes the target, prompt, and output length on both
/// sides. Different random-number consumption means sampled text need not be
/// identical; distribution-preservation tests are separate from work counts.
pub fn compare_tree_decoding(
  prefix : Array[Int],
  draft : (Array[Int]) -> Result[Array[Double], String],
  target : (Array[Int]) -> Result[Array[Double], String],
  config : TreeExperimentConfig,
) -> Result[TreeExperiment, TreeExperimentError] {
  let baseline = match decode_baseline_from_model(prefix, target, config) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  let speculative = match
    decode_tree_from_model(prefix, draft, target, config) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  Ok({ prompt: prefix.copy(), config, baseline, speculative })
}

///|
pub fn TreeExperiment::render(self : TreeExperiment) -> String {
  "TreeSpec context-dependent experiment\nseed=" +
  self.config.seed.to_string() +
  " width=" +
  self.config.width.to_string() +
  " depth=" +
  self.config.depth.to_string() +
  " node_budget=" +
  self.config.node_budget.to_string() +
  "\nbaseline_output_tokens=" +
  self.baseline.generated.length().to_string() +
  "\ntree_output_tokens=" +
  self.speculative.generated.length().to_string() +
  "\nbaseline_actual_target_queries=" +
  self.baseline.target_queries.to_string() +
  "\ntree_actual_target_queries=" +
  self.speculative.target_queries.to_string() +
  "\ntree_actual_draft_queries=" +
  self.speculative.draft_queries.to_string() +
  "\nbaseline_logical_target_batches=" +
  self.baseline.logical_target_batches.to_string() +
  "\ntree_logical_target_batches=" +
  self.speculative.logical_target_batches.to_string() +
  "\ncandidate_nodes=" +
  self.speculative.candidate_nodes.to_string() +
  "\naccepted_nodes=" +
  self.speculative.accepted_nodes.to_string() +
  "\nBackend: sequential synthetic logits; logical batches are hypothetical.\n" +
  "No hardware speedup is measured.\n"
}

///|
/// Small context-sensitive fixtures; no trained weights or external data.
pub fn run_tree_experiment_demo() -> Result[String, TreeExperimentError] {
  let config = match TreeExperimentConfig::new(2, 3, 14, 32, 20260905) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  let target = fn(context : Array[Int]) -> Result[Array[Double], String] {
    let last = if context.is_empty() {
      0
    } else {
      context[context.length() - 1]
    }
    Ok(if last == 0 { [0.0, 2.0, -1.0] } else { [1.0, -1.0, 0.5] })
  }
  let draft = fn(context : Array[Int]) -> Result[Array[Double], String] {
    let last = if context.is_empty() {
      0
    } else {
      context[context.length() - 1]
    }
    Ok(if last == 0 { [0.0, 1.0, -0.5] } else { [0.8, 0.0, 0.3] })
  }
  match compare_tree_decoding([0], draft, target, config) {
    Ok(result) => Ok(result.render())
    Err(error) => Err(error)
  }
}

///|
/// Runnable batch-backend companion to `run_tree_experiment_demo`. It uses the
/// same deterministic logits but exposes the provider-level request reduction
/// that an embedding inference runtime can turn into parallel work.
pub fn run_tree_batch_experiment_demo() -> Result[String, TreeExperimentError] {
  let config = match TreeExperimentConfig::new(2, 3, 14, 32, 20260905) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  let target = fn(
    contexts : Array[Array[Int]],
  ) -> Result[Array[Array[Double]], String] {
    let rows : Array[Array[Double]] = []
    for context in contexts {
      let last = if context.is_empty() {
        0
      } else {
        context[context.length() - 1]
      }
      rows.push(if last == 0 { [0.0, 2.0, -1.0] } else { [1.0, -1.0, 0.5] })
    }
    Ok(rows)
  }
  let draft = fn(
    contexts : Array[Array[Int]],
  ) -> Result[Array[Array[Double]], String] {
    let rows : Array[Array[Double]] = []
    for context in contexts {
      let last = if context.is_empty() {
        0
      } else {
        context[context.length() - 1]
      }
      rows.push(if last == 0 { [0.0, 1.0, -0.5] } else { [0.8, 0.0, 0.3] })
    }
    Ok(rows)
  }
  match decode_tree_from_batch_models([0], draft, target, config) {
    Ok(run) =>
      Ok(
        "TreeSpec batched tree experiment\noutput_tokens=" +
        run.generated.length().to_string() +
        "\nactual_target_batch_requests=" +
        run.target_queries.to_string() +
        "\nactual_draft_batch_requests=" +
        run.draft_queries.to_string() +
        "\nlogical_tree_rounds=" +
        run.logical_target_batches.to_string() +
        "\ncandidate_nodes=" +
        run.candidate_nodes.to_string() +
        "\n",
      )
    Err(error) => Err(error)
  }
}