object Expressions {
  import Exceptions._

  sealed abstract class ExprNodeLabel
  case class INT() extends ExprNodeLabel
  case class BOOL() extends ExprNodeLabel

  sealed abstract class ExprEdgeLabel
  case object add extends ExprEdgeLabel with EdgeLabel.Binary[INT, INT, INT]
  case object sub extends ExprEdgeLabel with EdgeLabel.Binary[INT, INT, INT]
  case object mul extends ExprEdgeLabel with EdgeLabel.Binary[INT, INT, INT]
  case object div extends ExprEdgeLabel with EdgeLabel.Binary[INT, INT, INT]
  case object lt extends ExprEdgeLabel with EdgeLabel.Binary[INT, INT, BOOL]
  case object gt extends ExprEdgeLabel with EdgeLabel.Binary[INT, INT, BOOL]
  case object equ extends ExprEdgeLabel with EdgeLabel.Binary[INT, INT, BOOL]
  case object and extends ExprEdgeLabel with EdgeLabel.Binary[BOOL, BOOL, BOOL]
  case object or extends ExprEdgeLabel with EdgeLabel.Binary[BOOL, BOOL, BOOL]

  type ExpressionCodeGraph = CodeGraph[ExprNodeLabel, ExprEdgeLabel]

  def evaluate(cg: ExpressionCodeGraph, inputs: Seq[Any]): Seq[Any] = {
    import CodeGraphOps._
    var env: Map[CodeGraph.NodeKey, Any] = cg.inputs.zip(inputs).toMap
    for (edgeKey <- cg.depthFirstOrder()) {
      val edge = cg.edge(edgeKey).head
      def args = edge.args.map(x => env.getOrElse(x, throw InvalidNodeKeyException(x)))
      def results = edge.label match {
        case `add` => Seq(args(0).asInstanceOf[Integer] + args(1).asInstanceOf[Integer])
        case `sub` => Seq(args(0).asInstanceOf[Integer] - args(1).asInstanceOf[Integer])
        case `mul` => Seq(args(0).asInstanceOf[Integer] * args(1).asInstanceOf[Integer])
        case `div` => Seq(args(0).asInstanceOf[Integer] / args(1).asInstanceOf[Integer])
        case `lt` => Seq(args(0).asInstanceOf[Integer] < args(1).asInstanceOf[Integer])
        case `gt` => Seq(args(0).asInstanceOf[Integer] > args(1).asInstanceOf[Integer])
        case `equ` => Seq(args(0).asInstanceOf[Integer] == args(1).asInstanceOf[Integer])
        case `and` => Seq(args(0).asInstanceOf[Boolean] && args(1).asInstanceOf[Boolean])
        case `or` => Seq(args(0).asInstanceOf[Boolean] || args(1).asInstanceOf[Boolean])
      }
      env = env ++ edge.results.zip(results).toMap
    }
    cg.outputs.map(x => env.getOrElse(x, throw InvalidNodeKeyException(x)))
  }

  object simpleExprCodeGraph extends CodeGraphBuilder[ExprNodeLabel, ExprEdgeLabel] {
    val (x, y) = input(INT, INT)

    val sum = add(x, y)
    val prod = mul(x, y)
    val prodMinusSum = sub(prod, sum)

    output(prodMinusSum)
  }
}