///| Decision traces make experiments explainable. They deliberately record

///| probabilities and random thresholds, not opaque model internals, so a

///|
/// trace can be inspected or serialized by a host application.
pub enum DecisionKind {
  Accepted
  Rejected
} derive(Eq, Debug)

///|
pub fn DecisionKind::label(self : DecisionKind) -> String {
  match self {
    Accepted => "accepted"
    Rejected => "rejected"
  }
}

///|
pub struct DecisionTrace {
  round : Int
  node_id : Int?
  token : Int
  kind : DecisionKind
  target_probability : Double
  draft_probability : Double
  threshold : Double
}

///|
pub fn DecisionTrace::summary(self : DecisionTrace) -> String {
  let node = match self.node_id {
    Some(id) => id.to_string()
    None => "path"
  }
  node + " token=" + self.token.to_string() + " kind=" + self.kind.label()
}

///|
pub struct DecodeTrace {
  entries : Array[DecisionTrace]
}

///|
pub fn DecodeTrace::empty() -> DecodeTrace {
  { entries: [] }
}

///|
pub fn DecodeTrace::push(self : DecodeTrace, entry : DecisionTrace) -> Unit {
  self.entries.push(entry)
}

///|
pub fn DecodeTrace::accepted_count(self : DecodeTrace) -> Int {
  let mut count = 0
  for entry in self.entries {
    if entry.kind == Accepted {
      count = count + 1
    }
  }
  count
}

///|
pub fn DecodeTrace::rejected_count(self : DecodeTrace) -> Int {
  let mut count = 0
  for entry in self.entries {
    if entry.kind == Rejected {
      count = count + 1
    }
  }
  count
}

///|
pub fn DecodeTrace::render(self : DecodeTrace) -> String {
  let mut text = "decision trace:\n"
  for entry in self.entries {
    text = text +
      "  - " +
      entry.summary() +
      " target=" +
      entry.target_probability.to_string() +
      " draft=" +
      entry.draft_probability.to_string() +
      " threshold=" +
      entry.threshold.to_string() +
      "\n"
  }
  text
}

///|
pub fn trace_single_path(
  round : Int,
  proposal : DraftProposal,
  target_distributions : Array[Array[Double]],
  accept_uniforms : Array[Double],
) -> Result[DecodeTrace, VerifyError] {
  match
    validate_verification_inputs(
      proposal, target_distributions, accept_uniforms,
    ) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  if target_distributions.length() != proposal.length() {
    return Err(TargetLengthMismatch)
  }
  if accept_uniforms.length() != proposal.length() {
    return Err(RandomLengthMismatch)
  }
  let trace = DecodeTrace::empty()
  for index in 0.. value
      Err(_) => return Err(ProbabilityFailure(index))
    }
    if accept_uniforms[index] < acceptance {
      trace.push({
        round,
        node_id: None,
        token: draft.token,
        kind: Accepted,
        target_probability,
        draft_probability,
        threshold: accept_uniforms[index],
      })
    } else {
      trace.push({
        round,
        node_id: None,
        token: draft.token,
        kind: Rejected,
        target_probability,
        draft_probability,
        threshold: accept_uniforms[index],
      })
      return Ok(trace)
    }
  }
  Ok(trace)
}