Encoding.scala

package de.tum.workflows

import scalax.collection.Graph
import scalax.collection.GraphPredef._
import scalax.collection.GraphEdge._
import scalax.collection.edge.LDiEdge
import scalax.collection.edge.Implicits._
import blocks._
import de.tum.workflows.foltl.FOLTL
import de.tum.workflows.foltl.FOLTL._
import de.tum.workflows.foltl.FormulaFunctions._
import com.typesafe.scalalogging._
import scalax.collection.edge.LBase.LEdgeImplicits
import scalax.collection.io.dot._
import implicits._
import Implicits._
import de.tum.workflows.toz3.InvProperties
import scalax.collection.GraphTraversal.{BreadthFirst, DepthFirst}

object Encoding extends LazyLogging {

  type WorkflowGraph = Graph[Int, LDiEdge]

  // Naming
  // TODO: make sure there are no clashes
  def nodeVar(n: Any) = Fun("n" + n, List())
  def choiceVar(n: Any) = "choice" + n

  def makeblock(b: Block, n1: Int, n2: Int, makeNode: (() => Int), makeId: (() => Block)): (List[LDiEdge[Int]]) = {

    b match {
      // single edge
      case b: SimpleBlock => {
        List((n1 ~+> n2)(b))
      }

      case Loop(steps) => {
        // add empty statement if loop ends with loop
        val adjsteps = steps.last match {
          case _: Loop  => steps :+ makeId()
          case _: Block => steps
        }
        // looping with n-1 nodes around n1
        val newnodes = n1 +: (0 until adjsteps.size - 1).toList.map(_ => makeNode()) :+ n1

        // n1 -> n2, n2 -> n3, ...
        val edgenodes = newnodes.sliding(2).toList
        val edges = for ((step, List(start, end)) <- adjsteps.zip(edgenodes)) yield {
          makeblock(step, start, end, makeNode, makeId)
        }
        edges.flatten :+ (n1 ~+> n2)(makeId())
      }

      case NondetChoice(l1, l2) => {
        // make sure that at least one of the two branches has one node, so they will never build 2 edges between the same node
        val newl1 = if (l1.length == l2.length && l1.length == 1) {
          l1 :+ makeId()
        } else l1

        val newnodes1 = n1 +: (0 until newl1.size - 1).toList.map(_ => makeNode()) :+ n2
        val edgenodes1 = newnodes1.sliding(2).toList
        val edges1 = for ((step, List(start, end)) <- newl1.zip(edgenodes1)) yield {
          makeblock(step, start, end, makeNode, makeId)
        }
        val newnodes2 = n1 +: (0 until l2.size - 1).toList.map(_ => makeNode()) :+ n2
        val edgenodes2 = newnodes2.sliding(2).toList
        val edges2 = for ((step, List(start, end)) <- l2.zip(edgenodes2)) yield {
          makeblock(step, start, end, makeNode, makeId)
        }
        edges1.flatten ++ edges2.flatten
      }
    }
  }

  def toGraph(w: Workflow): WorkflowGraph = {

    var n1 = 0

    def newnode() = {
      n1 = n1 + 1
      n1
    }
    def makeid() = ForallBlock(List(), List())

    val edges = (for ((step, i) <- w.steps.zipWithIndex) yield {
      if (i == w.steps.size - 1) {
        val n = n1
        val n2 = newnode()
        val edges = makeblock(step, n, n2, newnode, makeid)
        edges :+ (n2 ~+> n2)(makeid())
      } else {
        makeblock(step, n1, newnode(), newnode, makeid)
      }
    }).flatten
    val nodes = (0 until n1).toList

    Graph.from(nodes, edges)
  }

  def encodeGraph(g: WorkflowGraph) = {
    val allnodes = for (n <- g.nodes.toList) yield {
      val list = n.outgoing.toList.map(e => nodeVar(e.to))
      Implies(nodeVar(n), Next(Or.make(list: _*)))
    }
    Globally(And.make(allnodes))
  }

  def encodeSanity(g: WorkflowGraph) = {
    val negs = for (n <- g.nodes.toList; n2 <- g.nodes.toList if n < n2) yield {
      Neg(And(nodeVar(n), nodeVar(n2)))
    }
    val nodes = for (n <- g.nodes.toList if !(n == 0)) yield {
      Neg(nodeVar(n))
    }
    And.make(nodeVar(0), And.make(nodes), Globally(And.make(negs)))
  }

  def encodeStmt(s: Statement, pred: Option[Formula], bind: List[Var]): (Formula, String) = {
    //    var guard = s.guard
    //    for (p <- pred) { // if pred defined
    //      guard = And(guard, p)
    //    }
    val list = List(s.guard) ++ pred.toList
    val guard = And.make(list)

    // omitted variables
    val omitted = (bind toSet) -- (s.tuple toSet)

    val t = s match {
      case Add(_, f, tup) if omitted.isEmpty    => Or(Fun(f, tup), guard)
      case Add(_, f, tup)                       => Or(Fun(f, tup), Exists(omitted toList, guard))
      case Remove(_, f, tup) if omitted.isEmpty => And(Fun(f, tup), Neg(guard))
      case Remove(_, f, tup)                    => And(Fun(f, tup), Exists(omitted toList, Neg(guard)))
    }

    (Forall(s.tuple, Eq(Next(Fun(s.fun, s.tuple)), t)), s.fun)
  }

  def encodeStaying(target: Fun) = {
    Forall(target.params, Eq(Next(target), target))
  }

  def encodeSem(sig: Signature, g: Graph[Int, LDiEdge]) = {
    val impls = for (e <- g.edges.toList) yield {
      val choice = e.pred.map(n => Fun(n, e.agents))
      val bound = e.agents

      val res = for (s <- e.steps) yield {
        encodeStmt(s, choice, e.agents)
      }
      val updates = res.map(_._2)
      val staying = for (f <- sig.preds if !updates.contains(f.name)) yield {
        encodeStaying(f)
      }

      val terms = res.map(_._1) ++ staying
      Implies(And(nodeVar(e.from), Next(nodeVar(e.to))), And.make(terms))
    }
    Globally(And.make(impls))
  }

  def encodeInitial(sig: Signature, g: Graph[Int, LDiEdge]) = {
    val empties = for (f <- sig.preds toList) yield {
      Forall(f.params, Neg(f))
    }
    And.make(empties)
  }

  def sanecfg(g: Graph[Int, LDiEdge]) = {
    And(encodeGraph(g), encodeSanity(g))
  }

  def exec(sig: Signature, g: Graph[Int, LDiEdge]) = {
    And(encodeInitial(sig, g), encodeSem(sig, g))
  }

  def nodepredicates(g: Graph[Int, LDiEdge]) = {
    // toSet because the initial set is mutable
    g.nodes map nodeVar toSet
  }

  //  def toFOLTL(w: Workflow) = {
  //    val g = toGraph(w)
  //		logger.info("Graph of w:\n" + g)
  //
  //    val cfg = encodeGraph(g)
  //    logger.info(s"cfg(w): $cfg")
  //    val sanity = encodeSanity(g)
  //    logger.info(s"sanity(w): $sanity")
  //    val sem = encodeSem(w.sig, g)
  //    logger.info(s"sem(w): $sem")
  //    val init = encodeInitial(w.sig, g)
  //    logger.info(s"init(w): $init")
  //
  //    And.make(List(cfg, sanity, sem, init))
  //  }
  
  val MAXINVLENGTH = 800

  def toDot(g: WorkflowGraph, elaboration:Option[(InvProperties, Spec)] = None)(labels:Map[Int, String], edges:Set[g.EdgeT]) = {
    val root = DotRootGraph(
      directed = true,
      id = Some("Invariant Labelling"))

    def edgeTransformer(innerEdge: WorkflowGraph#EdgeT): Option[(DotGraph, DotEdgeStmt)] = innerEdge.edge match {
      case LDiEdge(source, target, label) => label match {
        case label: SimpleBlock => {
          val green = edges.exists(_ == innerEdge)
          val color = List(DotAttr("color", if (green) "green" else "red"))

          val elaborated = (elaboration map {
            case (prop, spec) => Preconditions.elaborate(label, spec, prop)
          }).getOrElse(List(label))
          val strlabel = elaborated.mkString("\n")

          val labelled = DotAttr("label", strlabel) :: color

          Some((
            root,
            DotEdgeStmt(
              source.toString,
              target.toString,
              labelled
            )
          ))
        }
      }
    }

    def nodeTransformer(innerNode: WorkflowGraph#NodeT): Option[(DotGraph, DotNodeStmt)] = {
      val str = labels.withDefault(d => s"$d").apply(innerNode.value)
      val label = if (str.length() > MAXINVLENGTH) str.substring(0, MAXINVLENGTH) + "..." + s"(${str.length()} characters)" else str
      Some((root, DotNodeStmt(innerNode.toString, List(DotAttr("label", "Node " + innerNode + ":\n" + label)))))
    }

    g.toDot(root, edgeTransformer, None, Some(nodeTransformer))
  }

  /**
    * Get an underapproximation of relation symbols that are sure to be untouched, i.e. empty
    *
    * @param w the graph
    */
  def untouchedMap(g:WorkflowGraph, sig:Signature): Map[g.NodeT, Set[String]] = {

    val sourceid = 0
    val sourcenode = g.get(sourceid)
    val nodes = sourcenode.withKind(BreadthFirst).filterNot(_ == sourcenode)

    // FIXME: this could 1. be more functional, 2. not just give up for loops
    var map:Map[g.NodeT, Set[String]] = Map().withDefault(_ => Set())

    // initially all untouched
    map = map.updated(g.get(sourceid), sig.preds.map(_.name))

    for (elem <- nodes) {
      // If all predecessors are known
      if (elem.diPredecessors.forall(map.contains(_))) {
        val union = elem.diPredecessors.flatMap(map(_))
        val touched = elem.incoming.flatMap(_.updating)
        map = map.updated(elem, union -- touched)
      }
    }

    map
  }

}