FormulaFunctions.scala

package de.tum.workflows.foltl

import FOLTL._
import com.typesafe.scalalogging.LazyLogging
import de.tum.workflows.blocks.Workflow
import scala.annotation.tailrec

object FormulaFunctions extends LazyLogging {

  def freeVars(t: Formula) = {
    def free(t: Formula, bound: Set[Var]): Set[Var] = {
      t match {
        // Extractors
        // Exists, Forall
        case Quantifier(_, ps, t) => free(t, bound ++ ps)
        case UnOp(make, t)        => free(t, bound)
        // And, Or, Eq, Implies
        case BinOp(_, t1, t2)     => free(t1, bound) ++ free(t2, bound)

        //      TODO: path is not a free variable anymore
        //        case Fun(_, Some(ind), ps) => (ps.toSet -- bound) + Var(ind)
        case Fun(_, _, ps)        => (ps.toSet -- bound)
        case v: Var if !bound(v)  => Set(v)
        case x                    => Set()
      }
    }
    free(t, Set())
  }

  def simplify(t: Formula): Formula = {
    val simp: PartialFunction[Formula, Formula] = {
      // Push negation inward
      case Neg(Or(t1, t2)) => And(Neg(t1), Neg(t2))
      case Neg(And(t1, t2)) => Or(Neg(t1), Neg(t2))
      case Neg(Globally(t)) => Finally(Neg(t))
      case Neg(Finally(t)) => Globally(Neg(t))
      case Neg(Next(t)) => Next(Neg(t))
      case Neg(Forall(vars, t)) => Exists(vars, Neg(t))
      case Neg(Exists(vars, t)) => Forall(vars, Neg(t))
      case Neg(Implies(t1, t2)) => And(t1, Neg(t2))

      // Simple laws
      case Neg(True) => False
      case Neg(False) => True
      case Neg(Neg(t)) => t
      case Or(True, t) => True
      case Or(t, True) => True
      case Or(False, t) => t
      case Or(t, False) => t
      case And(False, t) => False
      case And(t, False) => False
      case And(True, t) => t
      case And(t, True) => t
      case Implies(False, t) => True
      case Implies(True, t) => t
      case Or(Neg(t1), t2) if t1 == t2 => True
      case And(Neg(t1), t2) if t1 == t2 => False
      case Or(t1, Neg(t2)) if t1 == t2 => True
      case And(t1, Neg(t2)) if t1 == t2 => False

      // left-trees
      case And(And(t1, t2), t3) => And(t1, And(t2, t3))
      case Or(Or(t1, t2), t3) => Or(t1, Or(t2, t3))

      // equal things in tree
      case And(t1, t2) if t1 == t2 => t1
      case Or(t1, t2) if t1 == t2 => t1
      case And(t1, And(t2, t3)) if (t1 == t2) => And(t1, t3)
      case And(t1, And(t2, t3)) if (t1 == t3) => And(t1, t2)
      case Or(t1, Or(t2, t3)) if (t1 == t2) => Or(t1, t3)
      case Or(t1, Or(t2, t3)) if (t1 == t3) => Or(t1, t2)

      // And Or
      case And(t1, Or(t2, t3)) if (t1 == t3) => t1
      case And(t1, Or(t2, t3)) if (t1 == t2) => t1
      case And(Or(t2, t3), t1) if (t1 == t3) => t1
      case And(Or(t2, t3), t1) if (t1 == t2) => t1

      //      case Or(t1, Neg(t2)) if t1 == t2        => True

      // Equivalence
      case Eq(t1, t2) if t1 == t2 => True

      // Double Temporals
      case Finally(Finally(t)) => Finally(t)
      case Globally(Globally(t)) => Globally(t)

      // Remove quantifiers if empty
      case Quantifier(_, xs, t) if xs.isEmpty => t
      case Quantifier(_, _, True) => True
      case Quantifier(_, _, False) => False

      // Quantifiers binding same things
      case And(Forall(vars1, t1), Forall(vars2, t2)) if (vars1 == vars2) => Forall(vars1, And(t1, t2))
      case Or(Forall(vars1, t1), Forall(vars2, t2)) if (vars1 == vars2) => Forall(vars1, Or(t1, t2))
      case And(Exists(vars1, t1), Exists(vars2, t2)) if (vars1 == vars2) => Exists(vars1, And(t1, t2))
      case Or(Exists(vars1, t1), Exists(vars2, t2)) if (vars1 == vars2) => Exists(vars1, Or(t1, t2))

      // Gobble up equal quantifiers
      case Forall(vars1, Forall(vars2, f)) => Forall(vars1 ++ vars2, f)
      case Exists(vars1, Exists(vars2, f)) => Exists(vars1 ++ vars2, f)

      // Remove variables from quantifiers if not used in the body: Caution - expensive
      //      case Quantifier(qmake, xs, t) if !(xs.toSet -- t.freeVars()).isEmpty =>
      //        qmake((xs.toSet.intersect(t.freeVars())).toList, t)
      //
      //      // Split off binops from quantifiers
      //      case Quantifier(qmake, xs, BinOp(make, t1, t2)) if (t1.freeVars().intersect(xs.toSet)).isEmpty =>
      //        make(t1, qmake(xs, t2))
      //      case Quantifier(qmake, xs, BinOp(make, t1, t2)) if (t2.freeVars().intersect(xs.toSet)).isEmpty =>
      //        make(qmake(xs, t1), t2)

      // Declass seen often enough
      case And(Or(f1, f2), Or(f3, f4)) if (f2 == f4 && f1 == Neg(f3).simplify()) => f2
    }

    //    val (changed, t1) = everywhereFP(simp, t)
    //    if (!changed) t1 else simplify(t1)
    everywhereBotUp(simp, t)
  }

  def pushQuantifiers(f: Formula): Formula = {

    f everywhere {
      case Forall(vars, BinOp(make, f1, f2)) if {
        val leftvars = f1.freeVars
        val rightvars = f2.freeVars
        val varset = vars.toSet
        val onlyleft = leftvars.intersect(varset) -- rightvars
        val onlyright = rightvars.intersect(varset) -- leftvars
        !onlyleft.isEmpty || !onlyright.isEmpty
      } => {
        val leftvars = f1.freeVars
        val rightvars = f2.freeVars
        val varset = vars.toSet
        val bothvars = leftvars.intersect(rightvars).intersect(varset)
        val onlyleft = leftvars.intersect(varset) -- rightvars
        val onlyright = rightvars.intersect(varset) -- leftvars

        Forall.make(
          bothvars.toList,
          make(
            pushQuantifiers(Forall.make(onlyleft.toList, f1)),
            pushQuantifiers(Forall.make(onlyright.toList, f2))))
      }
    }
  }

  def eliminateImplication(f: Formula): Formula = {
    f everywhere {
      case Implies(f1, f2) => Or(Neg(eliminateImplication(f1)), eliminateImplication(f2))
    }
  }

  def eliminateEq(f: Formula): Formula = {
    f everywhere {
      case Eq(f1, f2) => {
        val e1 = eliminateEq(f1)
        val e2 = eliminateEq(f2)
        And(Implies(e1, e2), Implies(e2, e1))
      }
    }
  }

  def everywhereFP(trans: PartialFunction[Formula, Formula], t: Formula): (Boolean, Formula) = {
    if (trans.isDefinedAt(t))
      (true, trans(t))
    else
      t match {
        // Extractors
        case Quantifier(make, ps, t1) => {
          val sub = everywhereFP(trans, t1)
          (sub._1, make(ps, sub._2))
        }
        case UnOp(make, t1) => {
          val sub = everywhereFP(trans, t1)
          (sub._1, make(sub._2))
        }
        case BinOp(make, t1, t2) => {
          val (c1, sub1) = everywhereFP(trans, t1)
          val (c2, sub2) = everywhereFP(trans, t2)
          (c1 || c2, make(sub1, sub2))
        }
        case x => (false, x)
      }
  }

  def everywhere(trans: PartialFunction[Formula, Formula], t: Formula): Formula = {
    if (trans.isDefinedAt(t))
      trans(t)
    else
      t match {
        // Extractors
        case Quantifier(make, ps, t) => make(ps, everywhere(trans, t))
        case UnOp(make, t1)          => make(everywhere(trans, t1))
        case BinOp(make, t1, t2)     => make(everywhere(trans, t1), everywhere(trans, t2))
        case x                       => x
      }
  }

  def everywhereBotUp(trans: PartialFunction[Formula, Formula], t: Formula): Formula = {
    // First apply trans to push negation inward etc
    if (trans.isDefinedAt(t))
      everywhereBotUp(trans, trans(t))
    else {
      val sub = t match {
        // Extractors
        case Quantifier(make, ps, t) => make(ps, everywhereBotUp(trans, t))
        case UnOp(make, t1)          => make(everywhereBotUp(trans, t1))
        case BinOp(make, t1, t2)     => make(everywhereBotUp(trans, t1), everywhereBotUp(trans, t2))
        case x                       => x
      }
      // Applicable after inner transformation?
      if (trans.isDefinedAt(sub))
        trans(sub)
      else
        sub
    }
  }

  def opsize(t: Formula): Int = {
    t match {
      // Extractors
      case Quantifier(_, ps, t) => 1 + t.opsize
      case UnOp(_, t)           => 1 + t.opsize
      case BinOp(_, t1, t2)     => 1 + t1.opsize + t2.opsize
      case x                    => 1
    }
  }

  def collect[T](combine: (T*) => T)(coll: PartialFunction[Formula, T], t: Formula): T = {
    if (coll.isDefinedAt(t))
      coll(t)
    else
      t match {
        // Extractors
        case Quantifier(make, ps, t) => combine(collect(combine)(coll, t))
        case UnOp(make, t)           => combine(collect(combine)(coll, t))
        case BinOp(make, t1, t2)     => combine(collect(combine)(coll, t1), collect(combine)(coll, t2))
        case x                       => combine()
      }
  }

  def collect[T](coll: PartialFunction[Formula, List[T]], t: Formula): List[T] = {
    if (coll.isDefinedAt(t))
      coll(t)
    else
      t match {
        // Extractors
        case Quantifier(make, ps, t) => collect(coll, t)
        case UnOp(make, t)           => collect(coll, t)
        case BinOp(make, t1, t2)     => List.concat(collect(coll, t1), collect(coll, t2))
        case x                       => List.empty
      }
  }

  def hasSubFormula(coll: PartialFunction[Formula, Boolean], t: Formula): Boolean = {
    def or(bs: Boolean*) = { bs.exists(identity) }
    collect(or)(coll, t)
  }

  //  def hasSubFormula(check: (Formula => Boolean), t: Formula): Boolean = {
  //    if (check(t)) {
  //      true
  //    } else {
  //      t match {
  //        // Extractors
  //        case Quantifier(make, ps, t) => hasSubFormula(check, t)
  //        case UnOp(make, t)           => hasSubFormula(check, t)
  //        case BinOp(make, t1, t2)     => hasSubFormula(check, t1) || hasSubFormula(check, t2)
  //        case x                       => false
  //      }
  //    }
  //  }

  def assumeEmpty(t: Formula, name: String) = {
    t.everywhere({
      case Fun(f, _, _) if f == name => False
    })
  }

  def assumeAllEmpty(f: Formula, names: List[String]) = {
    // Assume all workflow relations empty
    val assumers = names.map(p => (f: Formula) => assumeEmpty(f, p))
    val allempty = assumers.foldLeft(f)((f, a) => a(f))
    allempty.simplify()
  }

  def annotate(t: Formula, name: String) = {
    t.everywhere({
      case Fun(f, None, xs) => Fun(f, Some(name), xs)
    })
  }

  def annotateOverwrite(t: Formula, Name1: String, name: String) = {
    t.everywhere({
      case Fun(f, Some(Name1), xs) => Fun(f, Some(name), xs)
    })
  }

  def annotate(t: Formula, name: String, ignore: Set[String]) = {
    t.everywhere({
      case Fun(f, None, xs) if (!(ignore contains f)) => Fun(f, Some(name), xs)
    })
  }

  def checkSanity(t: Formula) = {
    // TODO add more things, f.e. existentials
    // T1, T2 can appear freely
    val frees = t.freeVars -- Set(Var(Properties.T1), Var(Properties.T2))
    if (frees.nonEmpty) {
      logger.warn(s"Sanity check failed: $frees appear free in the property.")
    }

    frees.isEmpty
  }

  def removeOTQuantifiers(f: Formula): Formula = {
    f.everywhere({
      case ForallOtherThan(vars, otherthan, fp) => {
        val eqs = for (x <- vars; y <- otherthan) yield {
          Neg(Eq(x, y))
        }
        Forall(vars, Implies(And.make(eqs), fp.removeOTQuantifiers()))
      }
    })
  }

  // f[newvars/vars]
  def parallelRename(f: Formula, vars: List[Var], newvars: List[Var]): Formula = {
    val repls = vars zip newvars toMap

    parallelRename(f, repls)
  }

  def parallelRename(f: Formula, repls: Map[Var, Var]): Formula = {
    f everywhere {
      case v: Var if repls.contains(v) => repls(v)
      case Fun(f, ind, params)         => Fun(f, ind, params map (p => repls.getOrElse(p, p)))
      case Quantifier(make, xs, f) => {
        val rep = repls -- (xs)
        make(xs, parallelRename(f, rep))
      }
    }
  }

  /**
   * Check if there exist quantifiers binding the same variable more than once
   */
  def hasDuplicateBindings(f: Formula) = {
    def boundNames(f: Formula): List[Var] = {
      f collect {
        case Quantifier(make, xs, sub) => xs ++: boundNames(sub)
      }
    }
    boundNames(f).groupBy(n => n).exists(_._2.length > 1)
  }

  def stripQuantifiers(f: Formula): Formula = {
    f everywhere {
      case Quantifier(_, _, fsub) => stripQuantifiers(fsub)
    }
  }

  def collectQuantifiers(f: Formula): List[(Boolean, List[Var])] = {
    f collect {
      case Exists(xs, fsub) => (true, xs) :: collectQuantifiers(fsub)
      case Forall(xs, fsub) => (false, xs) :: collectQuantifiers(fsub)
      case BinOp(make, f1, f2) => {
        val q1 = collectQuantifiers(f1)
        val q2 = collectQuantifiers(f2)

        // Merge while trying to stay in BS, if possible
        // TODO: could be improved to minimize quantifier alternations)
        // TODO: merge universals with same type by renaming
        val ex1 = q1.takeWhile(_._1)
        val rest1 = q1.drop(ex1.size)

        // filter out same bound variables (TODO: could be improved if different names)
        val existbound = ex1.flatMap(_._2)
        val univbound = rest1.flatMap(_._2)

        val exset = existbound.toSet
        val univset = univbound.toSet

        val ex2 = q2.takeWhile(_._1).map(q =>
          (q._1, q._2.filterNot(exset.contains)))
        val rest2 = q2.drop(ex2.size).map(q =>
          (q._1, q._2.filterNot(univset.contains)))
        ex1 ++ ex2 ++ rest1 ++ rest2
      }
    }
  }

  def checkBindings(f: Formula) = {
    def checkBindingsSub(f: Formula, bound: Set[Var]): List[Var] = {
      f collect {
        case Quantifier(make, xs, inner) => {
          val conflicts = xs.filter(bound.contains)
          conflicts ++ checkBindingsSub(inner, bound ++ xs)
        }
      }
    }
    checkBindingsSub(f, Set()).isEmpty
  }

  def isBS(f: Formula) = {
    val quantprefix = collectQuantifiers(f.toNegNf)

    if (quantprefix.length < 2) true else {
      // has no forall exists in quantifier prefix
      !quantprefix.sliding(2).exists(list => !list.head._1 && list(1)._1)
    }
  }

  def isQFree(formula: Formula): Boolean = {
    !formula.hasSubFormula {
      case q: Quantifier => true
    }
  }

  def toNegNf(f: Formula) = {
    //    val bound = if (hasDuplicateBindings(f)) {
    //      //      logger.info(s"Formula binds variables more than once - renaming bindings: ${f.pretty}")
    //      FormulaFunctions.fixBinding(f, Set())._1
    //    } else f
    val noeq = eliminateEq(f)
    eliminateImplication(noeq).simplify()
  }

  def generateName(base: Var, bound: Set[Var]) = {

    val chars = base.name.toCharArray()
    val (basename, start) = if (chars.last.isDigit) {
      (base.name.dropRight(1), chars.last.toInt + 1)
    } else {
      (base.name, 1)
    }
    val stream = Stream.from(start)

    val droppedstream = stream.dropWhile(i => bound.contains(Var(basename + i, base.typ)))

    Var(basename + droppedstream.head, base.typ)
  }

  def fixBinding(f: Formula, bound: Set[Var]): (Formula, Set[Var]) = {
    f match {
      case UnOp(make, f1) => {
        val (newf, newbounds) = fixBinding(f1, bound)
        (make(newf), newbounds)
      }
      case BinOp(make, f1, f2) => {
        val (newf1, newbound) = fixBinding(f1, bound)
        val (newf2, newbound2) = fixBinding(f2, newbound)
        (make(newf1, newf2), newbound2)
      }
      case Quantifier(make, vars, fsub) => {
        val alreadybounds = vars.filter(bound.contains)
        val newnames = alreadybounds.map(generateName(_, bound))
        val renamed = fsub.parallelRename(alreadybounds, newnames)
        val binding = (vars.toSet -- alreadybounds) ++ newnames

        val (fixedrename, newbound) = fixBinding(renamed, bound ++ binding)
        val newf = make(binding.toList, fixedrename)
        (newf, newbound)
      }
      // Functions, Variables
      case f => (f, bound)
    }
  }

  // TODO: make sure no quantifiers in scope of others with same name
  def toPrenex(f: Formula): Formula = {
    val negnf = f.toNegNf
    //    val (simp, boundVars) = fixBinding(negnf, Set())

    // Simplify first to push negation inward
    val quants = collectQuantifiers(negnf)
    val stripped = stripQuantifiers(negnf)

    quants.foldRight(stripped)((quant, inner) => {
      if (quant._1) Exists(quant._2, inner) else Forall(quant._2, inner)
    }).simplify()
  }

  def toCNFClauses(f: Formula) = {
    @tailrec
    def toCNFSub(f: Formula): Formula = {
      val (changed, sub) = everywhereFP({
        case Or(And(f1, f2), f3) => And(Or(f1, f3), Or(f2, f3))
        case Or(f1, And(f2, f3)) => And(Or(f1, f2), Or(f1, f3))
      }, f)
      if (!changed) sub else toCNFSub(sub)
    }


    // FIXME: can we do it without prenex? may be expensive later on
    val normalized = f.toNegNf.simplify().toPrenex.simplify()
    logger.debug(s"Computing CNF of prenex formula of opsize ${normalized.opsize}")

    // normalized is now Quantifiers, then only And/Or with Negation inside
    // no simplification here
    val cnf1 = toCNFSub(normalized)
    val cnf2 = everywhereBotUp({
      case Or(And(f1, f2), f3) => And(Or(f1, f3), Or(f2, f3))
      case Or(f1, And(f2, f3)) => And(Or(f1, f2), Or(f1, f3))
    }, normalized)

    val quantifiers = FormulaFunctions.collectQuantifiers(cnf1)

    // Inspect clauses
    val clauses = FOTransformers.collectClauses(cnf1)

    // TODO: improve, don't sort by tostring
    val sorted = clauses.map(c => c.sortBy(_.toString())).sortBy(f => f.toString())

    logger.debug(s"Computing CNF of a formula with ${sorted.size} clauses")

    // Remove trivial clauses
    val notrivials = sorted.filterNot(c => {
      c.exists(f => c.contains(Neg(f)))
    })

    // TODO improve runtime
    val simped = notrivials.filterNot(c => notrivials.exists(c2 => {
      // check if c is superset of c2
      //      c != c2 && c.startsWith(c2)
      c != c2 && c.forall(f => c2.contains(f))
    }))
    logger.debug(s"After simplification ${simped.size} unsubsumed clauses remaining.")

    (quantifiers, simped)
  }

  def toCNF(f: Formula) = {
    val (quantifiers, clauses) = toCNFClauses(f)
    val form = And.make(clauses.map(Or.make))
    val newform = quantifiers.foldRight(form)((q, inner) => if (q._1) Exists(q._2, inner) else Forall(q._2, inner))
    newform
  }

  /**
   * Resolves a universally quantified predicate
   * Blows up the formula linearly due to CNF
   */
  def eliminatePredicate(f: Formula, name: String) = {
    def listFunsWithName(f: Formula, name: String) = {
      f.collect({
        case Fun(n, ind, vars) if n == name      => List((Some(Fun(n, ind, vars)), None))
        case Neg(Fun(n, ind, vars)) if n == name => List((None, Some(Neg(Fun(n, ind, vars)))))
      })
    }

    val norm = f.toCNF

    val removed = norm.everywhere({
      // find top-level ORs in CNF
      case inner: Or => {
        val list = listFunsWithName(inner, name)

        val positives = list.flatMap(_._1)
        val negatives = list.flatMap(_._2)

        val isSuperFluous = positives.exists { f1 =>
          negatives.contains(Neg(f1))
        }

        if (isSuperFluous) {
          // Remove clause
          True
        } else {
          // Remove predicate both in positive and negative
          inner.everywhere({
            case Neg(Fun(n, ind, vars)) if n == name => False
            case Fun(n, ind, vars) if n == name      => False
          })
        }
      }
    })

    removed.simplify()

  }
}