///| Deterministic offline simulator. It exercises the exact verifier without

///| model weights or network access, and is intentionally suitable for tests,

///|
/// CI benchmarks, and reproducible documentation examples.
pub enum SimulationError {
  EmptySchedule
  VocabularyMismatch
  ProbabilityError
  VerificationError
} derive(Eq, Debug)

///|
pub struct SimulationRound {
  draft_logits : Array[Array[Double]]
  target_logits : Array[Array[Double]]
  draft_uniforms : Array[Double]
  accept_uniforms : Array[Double]
  fallback_uniforms : Array[Double]
}

///|
pub struct SimulationResult {
  generated : Array[Int]
  metrics : DecodeMetrics
  rounds : Int
}

///|
pub fn argmax(values : Array[Double]) -> Result[Int, SimulationError] {
  if values.length() == 0 {
    return Err(EmptySchedule)
  }
  let mut best_index = 0
  let mut best_value = values[0]
  for index in 1.. best_value {
      best_index = index
      best_value = values[index]
    }
  }
  Ok(best_index)
}

///|
fn distributions_from_logits(
  logits : Array[Array[Double]],
) -> Result[Array[Array[Double]], SimulationError] {
  let output : Array[Array[Double]] = []
  for values in logits {
    match softmax(values) {
      Ok(distribution) => output.push(distribution)
      Err(_) => return Err(ProbabilityError)
    }
  }
  Ok(output)
}

///|
fn proposal_from_logits(
  prefix : Array[Int],
  logits : Array[Array[Double]],
  uniforms : Array[Double],
) -> Result[DraftProposal, SimulationError] {
  if uniforms.length() != logits.length() {
    return Err(ProbabilityError)
  }
  let distributions = match distributions_from_logits(logits) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  let tokens : Array[Int] = []
  for i in 0.. tokens.push(token)
      Err(_) => return Err(ProbabilityError)
    }
  }
  match make_proposal(prefix, tokens, distributions) {
    Ok(proposal) => Ok(proposal)
    Err(_) => Err(ProbabilityError)
  }
}

///|
pub fn run_round(
  prefix : Array[Int],
  round : SimulationRound,
) -> Result[VerifyResult, SimulationError] {
  if round.draft_logits.length() == 0 {
    return Err(EmptySchedule)
  }
  if round.draft_logits.length() != round.target_logits.length() {
    return Err(VocabularyMismatch)
  }
  let proposal = match
    proposal_from_logits(prefix, round.draft_logits, round.draft_uniforms) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  let targets = match distributions_from_logits(round.target_logits) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  match
    verify_proposal(
      proposal,
      targets,
      round.accept_uniforms,
      round.fallback_uniforms,
    ) {
    Ok(value) => Ok(value)
    Err(_) => Err(VerificationError)
  }
}

///|
pub fn simulate(
  prefix : Array[Int],
  schedule : Array[SimulationRound],
) -> Result[SimulationResult, SimulationError] {
  if schedule.length() == 0 {
    return Err(EmptySchedule)
  }
  let generated : Array[Int] = []
  for token in prefix {
    generated.push(token)
  }
  let metrics = DecodeMetrics::empty()
  let mut rounds = 0
  for round in schedule {
    let proposal = match
      proposal_from_logits(generated, round.draft_logits, round.draft_uniforms) {
      Ok(value) => value
      Err(error) => return Err(error)
    }
    let targets = match distributions_from_logits(round.target_logits) {
      Ok(value) => value
      Err(error) => return Err(error)
    }
    let result = match
      verify_proposal(
        proposal,
        targets,
        round.accept_uniforms,
        round.fallback_uniforms,
      ) {
      Ok(value) => value
      Err(_) => return Err(VerificationError)
    }
    metrics.record(result, proposal)
    for token in result.emitted {
      generated.push(token)
    }
    rounds = rounds + 1
  }
  Ok({ generated, metrics, rounds })
}

///|
pub struct AdaptivePolicy {
  mut depth : Int
  min_depth : Int
  max_depth : Int
  target_acceptance : Double
}

///|
pub fn AdaptivePolicy::new(
  min_depth : Int,
  max_depth : Int,
  target_acceptance : Double,
) -> AdaptivePolicy {
  { depth: min_depth, min_depth, max_depth, target_acceptance }
}

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

///|
pub fn AdaptivePolicy::observe(
  self : AdaptivePolicy,
  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
  }
}