///|
pub struct ExplainStep {
  iter : Int
  rewrite : String
  root_before : Id
  rhs_id : Id
  subst : Map[String, Id]
} derive(Show)

///|
pub struct ExplainEdge {
  from : Id
  to : Id
  rewrite : String
  iter : Int
  subst : Map[String, Id]
} derive(Show)

///|
pub struct Explanation {
  steps : Array[ExplainStep]
  edges : Map[Id, Array[ExplainEdge]]
}

///|
pub fn Explanation::new() -> Explanation {
  Explanation::{ steps: [], edges: Map::new() }
}

///|
pub fn Explanation::record(self : Explanation, step : ExplainStep) -> Unit {
  self.steps.push(step)
  let edge = ExplainEdge::{
    from: step.root_before,
    to: step.rhs_id,
    rewrite: step.rewrite,
    iter: step.iter,
    subst: step.subst,
  }
  self.edges.update(step.root_before, existing => match existing {
    Some(arr) => { let a = arr; a.push(edge); Some(a) }
    None => Some([edge])
  })
  // bidirectional for equivalence tracing
  self.edges.update(step.rhs_id, existing => match existing {
    Some(arr) => { let a = arr; a.push(edge); Some(a) }
    None => Some([edge])
  })
}

///|
pub fn Explanation::steps(self : Explanation) -> Array[ExplainStep] {
  self.steps
}

///|
/// Attempt to find a justification path between two e-class ids.
pub fn Explanation::explain_path(
  self : Explanation,
  start : Id,
  goal : Id,
) -> Array[ExplainEdge] {
  let start_id = start
  let goal_id = goal
  if start_id == goal_id {
    return []
  }
  let queue : Array[(Id, Array[ExplainEdge])] = [(start_id, [])]
  let seen : Map[Id, Bool] = Map::new()
  seen.set(start_id, true)
  loop () {
    _ =>
      match queue.pop() {
        None => break ()
        Some((node, path)) => {
          match self.edges.get(node) {
            None => continue ()
            Some(edges) =>
              for e in edges {
                let next = if e.from == node { e.to } else { e.from }
                if seen.contains(next) {
                  continue
                }
                let new_path = {
                  let res : Array[ExplainEdge] = Array::new()
                  res.append(path.copy().op_as_view())
                  res.push(e)
                  res
                }
                if next == goal_id {
                  return new_path
                }
                seen.set(next, true)
                queue.insert(0, (next, new_path))
              }
          }
        }
      }
  }
  []
}

///|
/// Pretty format a justification path between two e-classes, if any.
pub fn Explanation::format_path(
  self : Explanation,
  egraph : EGraph,
  start : Id,
  goal : Id,
) -> String {
  let path = self.explain_path(egraph.find(start), egraph.find(goal))
  if path.is_empty() && egraph.find(start) != egraph.find(goal) {
    return "no explanation".to_string()
  }
  let parts : Array[String] = Array::new()
  let mut cur = egraph.find(start)
  for edge in path {
    let next = if edge.from == cur { edge.to } else { edge.from }
    let subst_str = edge.subst.iter().map(pair => {
      let (k, v) = pair
      "\{k}=\{v}"
    }).join(",")
    parts.push("\{cur} --\{edge.rewrite}@iter\{edge.iter}[\{subst_str}]--> \{next}")
    cur = next
  }
  parts.join("\n")
}

///|
pub fn Explanation::has_path(
  self : Explanation,
  egraph : EGraph,
  a : Id,
  b : Id,
) -> Bool {
  let canon_a = egraph.find(a)
  let canon_b = egraph.find(b)
  if canon_a == canon_b {
    return true
  }
  !self.explain_path(canon_a, canon_b).is_empty()
}