///|
fn pattern_from_str(s : String) -> Pattern raise {
  parse_pattern(parse_sexpr(s))
}

///|
fn rewrite_from_str(name : String, lhs : String, rhs : String) -> Rewrite raise {
  Rewrite::new(name, pattern_from_str(lhs), pattern_from_str(rhs))
}

///|
fn rewrite_with_cond(
  name : String,
  lhs : String,
  rhs : String,
  cond : (EGraph, Map[String, Id]) -> Bool,
) -> Rewrite raise {
  Rewrite::with_condition(
    name,
    pattern_from_str(lhs),
    pattern_from_str(rhs),
    cond,
  )
}

///|
pub fn is_const(name : String) -> (EGraph, Map[String, Id]) -> Bool {
  (egraph, subst) => match subst.get(name) {
    Some(id) =>
      match egraph.data(id) {
        Some(payload) =>
          match payload.constant {
            Some(Value::Num(_)) => true
            _ => false
          }
        _ => false
      }
    None => false
  }
}

///|
pub fn is_sym(name : String) -> (EGraph, Map[String, Id]) -> Bool {
  (egraph, subst) => match subst.get(name) {
    Some(id) =>
      match egraph.class_for(id) {
        None => false
        Some(class) => {
          for idx in class.nodes {
            match egraph.nodes[idx].op {
              NodeOp::Symbol(_) => return true
              _ => ignore(())
            }
          }
          false
        }
      }
    None => false
  }
}

///|
pub fn is_not_zero(name : String) -> (EGraph, Map[String, Id]) -> Bool {
  (egraph, subst) => match subst.get(name) {
    Some(id) =>
      match egraph.data(id) {
        Some(payload) =>
          match payload.constant {
            Some(Value::Num(v)) => v != 0.0
            _ => true
          }
        _ => true
      }
    None => false
  }
}

///|
fn not_dep(
  var_name : String,
  expr_name : String,
) -> (EGraph, Map[String, Id]) -> Bool {
  (egraph, subst) => match (subst.get(var_name), subst.get(expr_name)) {
    (Some(v), Some(expr)) =>
      match egraph.data(expr) {
        Some(payload) => !payload.free.contains(egraph.find(v))
        None => false
      }
    _ => false
  }
}

///|
pub fn is_const_or_distinct(
  name : String,
  other : String,
) -> (EGraph, Map[String, Id]) -> Bool {
  (egraph, subst) => match (subst.get(name), subst.get(other)) {
    (Some(v), Some(o)) => {
      let v_root = egraph.find(v)
      let o_root = egraph.find(o)
      if v_root == o_root {
        return false
      }
      match egraph.data(v_root) {
        Some(payload) =>
          match payload.constant {
            Some(_) => true
            None =>
              match egraph.class_for(v_root) {
                Some(class) => {
                  for idx in class.nodes {
                    match egraph.nodes[idx].op {
                      NodeOp::Symbol(_) => return true
                      _ => ignore(())
                    }
                  }
                  false
                }
                None => false
              }
          }
        None => false
      }
    }
    _ => false
  }
}

///|
pub fn math_rules() -> Array[Rewrite] raise {
  [
    rewrite_from_str("comm-add", "(+ ?a ?b)", "(+ ?b ?a)"),
    rewrite_from_str("comm-mul", "(* ?a ?b)", "(* ?b ?a)"),
    rewrite_from_str("assoc-add", "(+ ?a (+ ?b ?c))", "(+ (+ ?a ?b) ?c)"),
    rewrite_from_str("assoc-add-rev", "(+ (+ ?a ?b) ?c)", "(+ ?a (+ ?b ?c))"),
    rewrite_from_str("assoc-mul", "(* ?a (* ?b ?c))", "(* (* ?a ?b) ?c)"),
    rewrite_from_str("sub-canon", "(- ?a ?b)", "(+ ?a (* -1 ?b))"),
    rewrite_with_cond(
      "div-canon",
      "(/ ?a ?b)",
      "(* ?a (pow ?b -1))",
      is_not_zero("b"),
    ),
    rewrite_from_str("zero-add", "(+ ?a 0)", "?a"),
    rewrite_from_str("zero-mul", "(* ?a 0)", "0"),
    rewrite_from_str("one-mul", "(* ?a 1)", "?a"),
    rewrite_from_str("add-zero", "?a", "(+ ?a 0)"),
    rewrite_from_str("mul-one", "?a", "(* ?a 1)"),
    rewrite_with_cond("cancel-div", "(/ ?a ?a)", "1", is_not_zero("a")),
    rewrite_from_str("cancel-sub", "(- ?a ?a)", "0"),
    rewrite_from_str(
      "distribute", "(* ?a (+ ?b ?c))", "(+ (* ?a ?b) (* ?a ?c))",
    ),
    rewrite_from_str("factor", "(+ (* ?a ?b) (* ?a ?c))", "(* ?a (+ ?b ?c))"),
    rewrite_from_str(
      "pow-mul", "(* (pow ?a ?b) (pow ?a ?c))", "(pow ?a (+ ?b ?c))",
    ),
    rewrite_with_cond("pow0", "(pow ?x 0)", "1", is_not_zero("x")),
    rewrite_from_str("pow1", "(pow ?x 1)", "?x"),
    rewrite_from_str("pow2", "(pow ?x 2)", "(* ?x ?x)"),
    rewrite_with_cond("pow-recip", "(pow ?x -1)", "(/ 1 ?x)", is_not_zero("x")),
    rewrite_with_cond("recip-mul-div", "(* ?x (/ 1 ?x))", "1", is_not_zero("x")),
    rewrite_with_cond("d-variable", "(d ?x ?x)", "1", is_sym("x")),
    rewrite_with_cond("d-constant", "(d ?x ?c)", "0", (egraph, subst) => is_sym(
        "x",
      )(egraph, subst) &&
      is_const_or_distinct("c", "x")(egraph, subst)),
    rewrite_with_cond("d-linear", "(d ?x (+ ?c (* ?k ?x)))", "?k", (
      egraph,
      subst,
    ) => {
      let ok_x = is_sym("x")(egraph, subst)
      let ok_c = not_dep("x", "c")(egraph, subst)
      let ok_k = not_dep("x", "k")(egraph, subst)
      ok_x && ok_c && ok_k
    }),
    rewrite_with_cond(
      "d-power-const-base",
      "(d ?x (pow ?x ?c))",
      "(* ?c (pow ?x (+ ?c -1)))",
      (egraph, subst) => is_sym("x")(egraph, subst) &&
        is_const("c")(egraph, subst),
    ),
    rewrite_from_str("d-add", "(d ?x (+ ?a ?b))", "(+ (d ?x ?a) (d ?x ?b))"),
    rewrite_from_str(
      "d-mul", "(d ?x (* ?a ?b))", "(+ (* ?a (d ?x ?b)) (* ?b (d ?x ?a)))",
    ),
    rewrite_from_str("d-sin", "(d ?x (sin ?x))", "(cos ?x)"),
    rewrite_from_str("d-cos", "(d ?x (cos ?x))", "(* -1 (sin ?x))"),
    rewrite_with_cond("d-ln", "(d ?x (ln ?x))", "(/ 1 ?x)", is_not_zero("x")),
    rewrite_with_cond(
      "d-power",
      "(d ?x (pow ?f ?g))",
      "(* (pow ?f ?g) (+ (* (d ?x ?f) (/ ?g ?f)) (* (d ?x ?g) (ln ?f))))",
      (egraph, subst) => is_not_zero("f")(egraph, subst) &&
        is_not_zero("g")(egraph, subst),
    ),
    rewrite_from_str("i-one", "(i 1 ?x)", "?x"),
    rewrite_with_cond(
      "i-power-const",
      "(i (pow ?x ?c) ?x)",
      "(/ (pow ?x (+ ?c 1)) (+ ?c 1))",
      is_const("c"),
    ),
    rewrite_from_str("i-cos", "(i (cos ?x) ?x)", "(sin ?x)"),
    rewrite_from_str("i-sin", "(i (sin ?x) ?x)", "(* -1 (cos ?x))"),
    rewrite_from_str("i-sum", "(i (+ ?f ?g) ?x)", "(+ (i ?f ?x) (i ?g ?x))"),
    rewrite_from_str("i-dif", "(i (- ?f ?g) ?x)", "(- (i ?f ?x) (i ?g ?x))"),
    rewrite_from_str(
      "i-parts", "(i (* ?a ?b) ?x)", "(- (* ?a (i ?b ?x)) (i (* (d ?x ?a) (i ?b ?x)) ?x))",
    ),
    rewrite_with_cond(
      "i-ln",
      "(i (ln ?x) ?x)",
      "(- (* ?x (ln ?x)) ?x)",
      is_not_zero("x"),
    ),
  ]
}