package de.tum.workflows.foltl

import FOLTL._
import com.typesafe.scalalogging.LazyLogging

object TermFunctions extends LazyLogging {

  def freeVars(t: Term) = {
    def free(t: Term, 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)

        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: Term): Term = {
    val simp: PartialFunction[Term, Term] = {
      // 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

      // 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(qmake, xs, t) if xs.isEmpty => t
      case Quantifier(qmake, xs, True)            => True
      case Quantifier(qmake, xs, False)           => False

      // Remove variables from quantifiers if not used in the body
      //      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)
    }

    val t1 = everywhere(simp, t)
    if (t == t1) t1 else simplify(t1)
  }

  def everywhere(trans: PartialFunction[Term, Term], t: Term): Term = {
    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 collect[T](combine:(T*) => T, empty:T)(coll: PartialFunction[Term, T], t: Term): T = {
    if (coll.isDefinedAt(t))
      coll(t)
    else
      t match {
        // Extractors
        case Quantifier(make, ps, t) => collect(combine, empty)(coll, t)
        case UnOp(make, t)           => collect(combine, empty)(coll, t)
        case BinOp(make, t1, t2)     => combine(collect(combine, empty)(coll, t1), collect(combine, empty)(coll, t2))
        case x                       => empty
      }
  }
  
   def collect[T](coll: PartialFunction[Term, List[T]], t: Term): 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 assumeEmpty(t: Term, name: String) = {
    t.everywhere({
      case Fun(f, _, _) if f == name => False
    })
  }

  def annotate(t: Term, name: String) = {
    t.everywhere({
      case Fun(f, None, xs) => Fun(f, Some(name), xs)
    })
  }
  
  def annotate(t: Term, name: String, ignore:Set[String]) = {
    t.everywhere({
      case Fun(f, None, xs) if (!(ignore contains f)) => Fun(f, Some(name), xs)
    })
  }
  
  def checkSanity(t: Term) = {
    // 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.isEmpty) {
      logger.warn(s"Sanity check failed: $frees appear free in the property.")
    }
    
    frees.isEmpty
  }
}