///| Comparison records and aggregate reports for deterministic decoding runs.

///| They intentionally store counters rather than wall-clock time: elapsed

///| time varies by device, whereas target calls and accepted draft tokens are

///|
/// reproducible algorithmic work measures.
pub enum ExperimentError {
  EmptyExperiment
  IncompatibleRecord(Int)
} derive(Eq, Debug)

///|
/// One named run under a fixed prompt and model fixture.
pub struct ExperimentRecord {
  name : String
  metrics : DecodeMetrics
  generated_tokens : Int
  prompt : Array[Int]
  target_model_id : String
}

///|
pub fn ExperimentRecord::from_result(
  name : String,
  target_model_id : String,
  prompt : Array[Int],
  result : SimulationResult,
) -> Result[ExperimentRecord, ExperimentError] {
  if result.generated.length() < prompt.length() || target_model_id.is_empty() {
    return Err(IncompatibleRecord(0))
  }
  for i in 0.. Int {
  self.metrics.target_batches
}

///|
pub fn ExperimentRecord::tokens_per_call(self : ExperimentRecord) -> Double {
  self.metrics.tokens_per_target_batch()
}

///|
/// Imported records carry caller-declared provenance. For internally matched
/// model execution prefer compare_tree_decoding, which runs both sides itself.
pub struct ExperimentPair {
  case_name : String
  baseline : ExperimentRecord
  speculative : ExperimentRecord
}

///|
pub fn ExperimentPair::new(
  case_name : String,
  baseline : ExperimentRecord,
  speculative : ExperimentRecord,
) -> Result[ExperimentPair, ExperimentError] {
  if baseline.generated_tokens != speculative.generated_tokens ||
    baseline.prompt != speculative.prompt ||
    baseline.target_model_id != speculative.target_model_id ||
    baseline.target_model_id.is_empty() ||
    baseline.metrics.emitted_tokens != baseline.generated_tokens ||
    speculative.metrics.emitted_tokens != speculative.generated_tokens {
    return Err(IncompatibleRecord(0))
  }
  Ok({ case_name, baseline, speculative })
}

///|
pub fn ExperimentPair::logical_batch_reduction(self : ExperimentPair) -> Double {
  if self.baseline.metrics.target_batches == 0 {
    0.0
  } else {
    1.0 -
    self.speculative.metrics.target_batches.to_double() /
    self.baseline.metrics.target_batches.to_double()
  }
}

///|
pub fn ExperimentPair::accepted_per_draft_token(
  self : ExperimentPair,
) -> Double {
  self.speculative.metrics.acceptance_rate()
}

///|
pub fn ExperimentPair::render(self : ExperimentPair) -> String {
  self.case_name +
  " baseline_calls=" +
  self.baseline.metrics.target_batches.to_string() +
  " speculative_calls=" +
  self.speculative.metrics.target_batches.to_string() +
  " reduction=" +
  self.logical_batch_reduction().to_string() +
  " acceptance=" +
  self.accepted_per_draft_token().to_string()
}

///|
/// Aggregate view of paired target-call savings and speculative acceptance.
pub struct ExperimentSummary {
  cases : Int
  reduction : SampleSummary
  acceptance : SampleSummary
  speculative_tokens_per_call : SampleSummary
}

///|
pub fn summarize_experiment(
  pairs : Array[ExperimentPair],
) -> Result[ExperimentSummary, ExperimentError] {
  if pairs.length() == 0 {
    return Err(EmptyExperiment)
  }
  let reductions : Array[Double] = []
  let acceptances : Array[Double] = []
  let throughput : Array[Double] = []
  for pair in pairs {
    reductions.push(pair.logical_batch_reduction())
    acceptances.push(pair.accepted_per_draft_token())
    throughput.push(pair.speculative.tokens_per_call())
  }
  let reduction = match summarize_sample(reductions) {
    Ok(value) => value
    Err(_) => return Err(EmptyExperiment)
  }
  let acceptance = match summarize_sample(acceptances) {
    Ok(value) => value
    Err(_) => return Err(EmptyExperiment)
  }
  let speculative_tokens_per_call = match summarize_sample(throughput) {
    Ok(value) => value
    Err(_) => return Err(EmptyExperiment)
  }
  Ok({
    cases: pairs.length(),
    reduction,
    acceptance,
    speculative_tokens_per_call,
  })
}

///|
pub fn ExperimentSummary::render_markdown(self : ExperimentSummary) -> String {
  "# TreeSpec experiment summary\n\n" +
  "cases=" +
  self.cases.to_string() +
  "\n\n" +
  self.reduction.render("logical_batch_reduction") +
  "\n" +
  self.acceptance.render("draft_acceptance") +
  "\n" +
  self.speculative_tokens_per_call.render("speculative_tokens_per_target_batch") +
  "\n"
}