package object algebraic {

  // normal AST
  trait Expr
  case class Const(x: Int) extends Expr
  case class Var(n: String) extends Expr
  case class Sum(a: Expr, b: Expr) extends Expr

  def eval(e: Expr, vars: Map[String, Int]): Int = {
    e match {
      case Const(x) => x
      case Sum(a, b) => eval(a, vars) + eval(b, vars)
      case Var(n) => vars(n)
    }
  }

  def format(e: Expr): String = {
    e match {
      case Const(x) => x.toString
      case Sum(a, b) => s"(${format(a)} + ${format(b)})"
      case Var(n) => n
    }
  }

  // convert (((5 + a) + (6 + b)) + 0) to (5 + (a + (6 + (b + 0))))
  def flatten(e: Expr): Expr = {
    e match {
      case v@Var(_) => v
      case c@Const(_) => c
      case Sum(a@Var(_), b) => Sum(a, flatten(b))
      case Sum(a@Const(_), b) => Sum(a, flatten(b))
      case Sum(Sum(a, b), c) => flatten(Sum(a, Sum(b, c)))
    }
  }

  def simplify(e: Expr): Expr = {
    def inner(e: Expr): Expr = {
      e match {
        case Sum(Var(a), Sum(Var(b), c)) =>
          if (a < b) Sum(Var(a), Sum(Var(b), inner(c)))
          else Sum(Var(b), Sum(Var(a), inner(c)))
        case Sum(Const(a), Sum(Var(b), c)) => Sum(Var(b), inner(Sum(Const(a), c)))
        case Sum(Var(a), Sum(Const(b), c)) => Sum(Var(a), inner(Sum(Const(b), c)))
        case Sum(Const(a), Sum(Const(b), c)) => inner(Sum(Const(a + b), c))
        case Sum(Const(a), Const(b)) => Const(a + b)
        case Sum(v@Var(_a), c@Const(_b)) => Sum(v, c)
        case Sum(c@Const(_a), v@Var(_b)) => Sum(v, c)
        case Sum(va@Var(a), vb@Var(b)) => if (a < b) Sum(va, vb) else Sum(vb, va)
        case v@Var(_) => v
        case c@Const(_) => c
      }
    }
    inner(flatten(e))
  }
}

package object visitor {

  // AST with visitor enabled
  trait Expr {
    def accept[A](op: Op[A]): A
  }
  case class Const(x: Int) extends Expr {
    def accept[A](op: Op[A]): A = op.apply(this)
  }
  case class Var(n: String) extends Expr {
    def accept[A](op: Op[A]): A = op.apply(this)
  }
  case class Sum(a: Expr, b: Expr) extends Expr {
    def accept[A](op: Op[A]): A = op.apply(this)
  }

  trait Op[A] {
    def apply(c: Const): A
    def apply(v: Var): A
    def apply(s: Sum): A
  }

  class Eval(vars: Map[String, Int]) extends Op[Int] {
    override def apply(c: Const): Int = c.x
    override def apply(v: Var): Int = vars(v.n)
    override def apply(s: Sum): Int = s.a.accept(this) + s.b.accept(this)
  }

  class Format extends Op[String] {
    override def apply(c: Const): String = c.x.toString
    override def apply(v: Var): String = v.n
    override def apply(s: Sum): String =
      s"(${s.a.accept(this)} + ${s.b.accept(this)})"
  }

  class Flatten extends Op[Expr] {
    override def apply(c: Const): Expr = c
    override def apply(v: Var): Expr = v
    override def apply(s: Sum): Expr = {
      s.a match {
        case _: Const => Sum(s.a, s.b.accept(this))
        case _: Var => Sum(s.a, s.b.accept(this))
        case t: Sum => Sum(t.a, Sum(t.b, s.b)).accept(this)
      }
    }
  }

  class Simplify extends Op[Expr] {
    var ex: Expr = _
    override def apply(c: Const): Expr = c
    override def apply(v: Var): Expr = v
    override def apply(s: Sum): Expr = {
      if (ex == null) {
        ex = s.accept(new Flatten)
      }
      ex match {
        case a: Const => s.b match {
          case b: Const => Const(a.x + b.x)
          case b: Var => Sum(b, a)
          case b: Sum => ???
          }
        case a: Var => s.b match {
          case b: Const => Sum(a, b)
          case b: Var => if (a.n < b.n) Sum(a, b) else Sum(b, a)
          case b: Sum => ???
        }
        case a: Sum => s.b match {
          case b: Const => ???
          case b: Var => ???
          case b: Sum => ???
        }
      }
    }
  }

  def eval(expr: Expr, vars: Map[String, Int]): Int = expr.accept(new Eval(vars))

  def format(expr: Expr): String = expr.accept(new Format)

  def flatten(expr: Expr): Expr = expr.accept(new Flatten)

  def simplify(expr: Expr): Expr = expr.accept(new Simplify)
}

object Test {

  import algebraic.{Const => ConstA, Var => VarA, Sum => SumA, eval => evalA, format => formatA, flatten => flattenA, simplify => simplifyA}
  import visitor.{Const => ConstV, Var => VarV, Sum => SumV, eval => evalV, format => formatV, flatten => flattenV, simplify => simplifyV}

  def main(args: Array[String]): Unit = {
    val m = Map("a" -> 1, "b" -> 2)

    val exa = SumA(SumA(SumA(ConstA(5), VarA("a")), SumA(ConstA(6), VarA("b"))), ConstA(0))
    println(s"algebraic eval on $m: ${evalA(exa, m)}")
    println(s"algebraic format: ${formatA(exa)}")
    println(s"algebraic flatten: ${formatA(flattenA(exa))}")
    println(s"algebraic simplify: ${formatA(simplifyA(exa))}")
    println()

    val exv = SumV(SumV(SumV(ConstV(5), VarV("a")), SumV(ConstV(6), VarV("b"))), ConstV(0))
    println(s"visitor eval on $m: ${evalV(exv, m)}")
    println(s"visitor format: ${formatV(exv)}")
    println(s"visitor flatten: ${formatV(flattenV(exv))}")
    println(s"visitor simplify: ${formatV(simplifyV(exv))}")
    println()

    //    val ex = Sum(
    //      Sum(Sum(Const(1),Const(2)),Sum(Const(3),Const(4))),
    //      Sum(Sum(Const(5),Const(6)),Sum(Const(7),Const(8)))
    //    )
  }
}