///|
pub enum StopReason {
  Saturated
  IterationLimit
  NodeLimit
  MatchLimit
} derive(Show, Eq)

///|
pub(all) enum Worklist {
  All
  Queue
  Backoff
  Greedy
  Recent
} derive(Show, Eq)

///|
pub struct RunConfig {
  iter_limit : Int
  node_limit : Int?
  match_limit : Int?
  worklist : Worklist
  allow_ematching_cycles : Bool
}

///|
pub fn RunConfig::default() -> RunConfig {
  RunConfig::{
    iter_limit: 5,
    node_limit: None,
    match_limit: None,
    worklist: Worklist::All,
    allow_ematching_cycles: true,
  }
}

///|
pub fn RunConfig::new(
  iter_limit : Int,
  node_limit? : Int? = None,
  match_limit? : Int? = None,
  worklist? : Worklist = Worklist::All,
  allow_ematching_cycles? : Bool = true,
) -> RunConfig {
  RunConfig::{ iter_limit, node_limit, match_limit, worklist, allow_ematching_cycles }
}

///|
/// Return ids reachable from the given root via enode children.
pub fn reachable_classes(egraph : EGraph, root : Id) -> Map[Id, Bool] {
  let seen : Map[Id, Bool] = Map::new()
  let stack : Array[Id] = [egraph.find(root)]
  loop () {
    _ =>
      match stack.pop() {
        None => break ()
        Some(id) => {
          let canon = egraph.find(id)
          if seen.contains(canon) {
            continue ()
          }
          seen.set(canon, true)
          match egraph.class_for(canon) {
            Some(class) =>
              for node_idx in class.nodes {
                let node = egraph.nodes[node_idx]
                for child in node.children {
                  stack.push(child)
                }
              }
            None => ()
          }
          continue ()
        }
      }
  }
  seen
}

///|
pub struct RunResult {
  egraph : EGraph
  root : Id
  iterations : Int
  stop_reason : StopReason
  applied : Map[String, Int]
  total_applied : Int
  reports : Array[IterationReport]
  explanation : Explanation
}

///|
pub struct IterationReport {
  iter : Int
  matched : Map[String, Int]
  applied : Map[String, Int]
  total_matched : Int
  total_applied : Int
  node_count : Int
  class_count : Int
  saturated : Bool
}

///|
pub struct Runner {
  rewrites : Array[Rewrite]
  mut queue : Array[Rewrite]
  egraph : EGraph
  mut root : Id
  config : RunConfig
  backoff : Map[String, Int]
  last_applied : Map[String, Int]
}

///|
pub fn Runner::new(
  expr : Expr,
  rewrites : Array[Rewrite],
  config? : RunConfig = RunConfig::default(),
) -> Runner {
  let egraph = EGraph::new()
  let root = egraph.add_expr(expr)
  egraph.rebuild()
  let queue = match config.worklist {
    Worklist::Queue => rewrites.copy()
    Worklist::All | Worklist::Backoff => []
    Worklist::Greedy | Worklist::Recent => []
  }
  egraph.set_allow_cycles(config.allow_ematching_cycles)
  Runner::{ rewrites, queue, egraph, root, config, backoff: Map::new(), last_applied: Map::new() }
}

///|
pub fn Runner::run(self : Runner) -> RunResult raise {
  let mut iter = 0
  let mut reason = StopReason::Saturated
  let stats : Map[String, Int] = Map::new()
  let mut total_applied = 0
  let reports : Array[IterationReport] = Array::new()
  let explanation = Explanation::new()
  loop () {
    _ => {
      let _ = match self.config.node_limit {
        Some(limit) if self.egraph.nodes.length() > limit => {
          reason = StopReason::NodeLimit
          break ()
        }
        _ => ()
      }
      if iter >= self.config.iter_limit {
        reason = StopReason::IterationLimit
        break ()
      }
      let mut changed = false
      let matched_map : Map[String, Int] = Map::new()
      let applied_map : Map[String, Int] = Map::new()
      let mut matched_total = 0
      let active : Array[Rewrite] = match self.config.worklist {
        Worklist::All => self.rewrites
        Worklist::Queue =>
          if self.queue.is_empty() {
            self.rewrites
          } else {
            let next = self.queue.remove(0)
            let remaining = self.queue
            self.queue = remaining
            self.queue.push(next)
            [next]
          }
        Worklist::Backoff => {
          let usable : Array[Rewrite] = Array::new()
          for rw in self.rewrites {
            let count = self.backoff.get(rw.name).unwrap_or(0)
            if count > 0 {
              self.backoff.set(rw.name, count - 1)
              continue
            }
            usable.push(rw)
          }
          usable
        }
        Worklist::Greedy => {
          let ordered : Array[Rewrite] = Array::new()
          let temp = self.rewrites
          let scores = self.last_applied
          let used : Map[String, Bool] = Map::new()
          loop () {
            _ => {
              let mut best : (Int, Int)? = None // (score, index)
              let mut best_rw : Rewrite? = None
              for i in 0.. {
                    best = Some((score, i))
                    best_rw = Some(rw)
                  }
                  Some((s, _)) =>
                    if score > s {
                      best = Some((score, i))
                      best_rw = Some(rw)
                    }
                }
              }
              match best_rw {
                None => break ()
                Some(rw) => {
                  ordered.push(rw)
                  used.set(rw.name, true)
                  continue ()
                }
              }
            }
          }
          ordered
        }
        Worklist::Recent => {
          let nonzero : Array[Rewrite] = self.rewrites.filter(rw => self.last_applied.get(rw.name).unwrap_or(0) > 0)
          if nonzero.is_empty() {
            self.rewrites
          } else {
            nonzero
          }
        }
      }
      for rw in active {
        let report = rw.apply_all_stats(self.egraph, on_applied=(m, rhs) => explanation.record(ExplainStep::{
            iter,
            rewrite: rw.name,
            root_before: self.egraph.find(m.root),
            rhs_id: self.egraph.find(rhs),
            subst: m.subst,
          },
        ))
        matched_total = matched_total + report.matched
        let applied = report.applied
        changed = changed || applied > 0
        if self.config.worklist == Worklist::Backoff {
          if applied > 0 {
            self.backoff.set(rw.name, 0)
          } else {
            let prev = self.backoff.get(rw.name).unwrap_or(0)
            let next = if prev >= 4 { 4 } else { prev + 1 }
            self.backoff.set(rw.name, next)
          }
        }
        if self.config.worklist == Worklist::Greedy ||
          self.config.worklist == Worklist::Recent {
          self.last_applied.set(rw.name, applied)
        }
        matched_map.update(report.name, existing => match existing {
          Some(v) => Some(v + report.matched)
          None => Some(report.matched)
        })
        if applied > 0 {
          total_applied = total_applied + applied
          stats.update(report.name, existing => match existing {
            Some(v) => Some(v + applied)
            None => Some(applied)
          })
          applied_map.update(report.name, existing => match existing {
            Some(v) => Some(v + applied)
            None => Some(applied)
          })
        }
        match self.config.match_limit {
          Some(limit) if total_applied >= limit => {
            reason = StopReason::MatchLimit
            break
          }
          _ => ignore(())
        }
      }
      if reason == StopReason::MatchLimit {
        self.egraph.rebuild()
        self.root = self.egraph.find(self.root)
        let class_count = self.egraph.class_ids().length()
        let node_count = self.egraph.nodes.length()
        reports.push(IterationReport::{
          iter,
          matched: matched_map,
          applied: applied_map,
          total_matched: matched_total,
          total_applied,
          node_count,
          class_count,
          saturated: !changed,
        })
        break ()
      }
      self.egraph.rebuild()
      self.root = self.egraph.find(self.root)
      self.egraph.set_allow_cycles(self.config.allow_ematching_cycles)
      let class_count = self.egraph.class_ids().length()
      let node_count = self.egraph.nodes.length()
      reports.push(IterationReport::{
        iter,
        matched: matched_map,
        applied: applied_map,
        total_matched: matched_total,
        total_applied,
        node_count,
        class_count,
        saturated: !changed,
      })
      iter = iter + 1
      let _ = match self.config.node_limit {
        Some(limit) if self.egraph.nodes.length() > limit => {
          reason = StopReason::NodeLimit
          break ()
        }
        _ => ()
      }
      if !changed {
        reason = StopReason::Saturated
        break ()
      }
      continue ()
    }
  }
  RunResult::{
    egraph: self.egraph,
    root: self.root,
    iterations: iter,
    stop_reason: reason,
    applied: stats,
    total_applied,
    reports,
    explanation,
  }
}

///|
pub fn run(
  expr : Expr,
  rewrites : Array[Rewrite],
  config? : RunConfig = RunConfig::default(),
) -> RunResult raise {
  Runner::new(expr, rewrites, config~).run()
}

///|
pub fn run_rewrites(
  expr : Expr,
  rewrites : Array[Rewrite],
  iter_limit? : Int = 5,
) -> (SimpleGraph, Id) raise {
  let result = run(
    expr,
    rewrites,
    config=RunConfig::new(iter_limit, node_limit=None),
  )
  (result.egraph, result.root)
}