
package uk.ac.man.cs.lethe.internal.fol.datatypes

import scala.collection.immutable.TreeSet
import scala.collection.immutable.SortedSet

abstract class Expression { 
  def contains(term: Term): Boolean
  def variables: Set[Variable]
}

//===========
//   Terms
//===========

abstract class Term extends Expression with Ordered[Term]{
  def subterms = Set(this)
}

case class Variable(name: String) extends Term { 
  override def toString = "v: "+name
  override def contains(term: Term) = term==this

  override def compare(that: Term) = that match { 
    case v: Variable => name.compare(v.name)
    case _ => -1
  }

  override def variables = Set(this)

  override def equals(other: Any) = other match { 
    case Variable(name2) => name==name2
    case _ => false
  }
}

object VariableBuilder {
  var counter = 0
  def newVariable(): Variable = {
    counter = counter + 1
    new Variable("_x"+counter)
  }

  def reinit() = 
    counter = 0
}

case class Constant(name: String) extends Term { 
  override def toString: String = "c: "+name
  override def contains(term: Term) = term==this

  override def compare(that: Term) = that match { 
    case v: Variable => 1
    case c: Constant => -name.compare(c.name)
    case _ =>  -1
  }

  override def variables = Set()

  override def equals(other: Any) = other match { 
    case Constant(name2) => name==name2
    case _ => false
  }
}
 
case class Function(name: String, args: List[Term]) extends Term {
  override def toString: String = name + args.mkString("(", ", ", ")")
  override def contains(term: Term) = this==term || args.exists(_.contains(term))
  override def subterms = args.map(_.subterms).toSet.flatten + this;

  override def compare(that: Term) = that match { 
    case Function(n2, a2) if n2==name => compare(args, a2)
    case Function(n2, _) => -name.compare(n2)
    case _ => 1
  }

  def compare(t1:List[Term], t2:List[Term]): Int = (t1,t2) match { 
    case (Nil, Nil) => 0
    case (Nil, _) => 1
    case (_, Nil) => -1
    case (a::_, b::_) if(a!=b) => a.compare(b)
    case (_::as, _::bs) => compare(as,bs)
  }

  override def variables = args.flatMap(_.variables).toSet
}

abstract class Predicate(_name: String, _args: List[Term]) extends Expression with Ordered[Predicate] {
  val args: List[Term] = _args
  val name: String = _name
  
  override def contains(term: Term) = _args.exists(_.contains(term))

  override def compare(that: Predicate) =
    if(that==null)
      1
    else if(name!=that.name) 
      name.compare(that.name) 
    else 
      compare(args, that.args)

  def compare(a1: List[Term], a2: List[Term]): Int = (a1, a2) match { 
    case(Nil, Nil) => 0
    case(Nil, _) => 1
    case(_, Nil) => -1
    case(a::_, b::_) if(a!=b) => a.compare(b)
    case(_::as, _::bs) => compare(as,bs)
  }
    
  override def variables = args.flatMap(_.variables).toSet

  override def equals(that: Any) = that match { 
    case p: Predicate => name==p.name && args==p.args
  }
}

case class GenericPredicate(_name: String, _args: List[Term]) extends Predicate(_name, _args) {
  override def toString: String = name + args.mkString("(", ", ", ")")

  override def equals(that: Any) = that match { 
    case p: GenericPredicate => name==p.name && args==p.args
  }

  override def hashCode = name.hashCode*31 + args.hashCode


}

case class IdentityPredicate(arg1: Term, arg2: Term) extends Predicate("=", List(arg1,arg2)) {
  override def toString: String = arg1.toString + " = " + arg2.toString

  override def equals(that: Any) = that match { 
    case IdentityPredicate(a,b) => (arg1==a && arg2==b) || (arg1==b && arg2==a)
    case _ => false
  }

  override def hashCode() = arg1.hashCode() + arg2.hashCode
  override def compare(that: Predicate) = that match { 
    case IdentityPredicate(arg1a, arg2a) => 
      val result = arg1.compare(arg1a)
      if(result==0)
	arg2.compare(arg2a)
      else
	result
    case _ => super.compare(that)
  }
}

abstract class Formula extends Expression { 
  def tautology = false
  def contradiction = false
  def signature: Set[String]
  def literals: Set[Literal]
}

// Only for testing purposes at the moment
case class SimpleProposition(name: String) extends Formula { 
  override def toString: String = name
  override def contains(term: Term) = false
  override def signature = Set(name)
  override def literals = Set()
  override def variables = Set()
}

case class Literal(positive: Boolean, predicate:Predicate) extends Formula with Ordered[Literal] {
  override def toString: String = predicate match {
    case IdentityPredicate(v1, v2) if !positive => v1.toString + " != " + v2.toString
    case _ if !positive => "¬"+predicate
    case _ => predicate.toString
  }

  override def signature = Set(predicate.name)

  override def contains(term: Term) = predicate.contains(term)

  override def compare(that: Literal) = { 
    var result = predicate.name.compare(that.predicate.name)
    if(result==0){ 
      if(positive && !that.positive)
	result = -1
      else if(!positive && that.positive)
	result = 1
      else 
	result = compare(predicate.args, that.predicate.args)
    }

    assert(result!=0 || this==that, "compare inconsistent with equals! "+toString+", "+that.toString)
    result
  }

  def compare(t1:List[Term], t2:List[Term]): Int = (t1,t2) match { 
    case (Nil, Nil) => 0
    case (Nil, _) => 1
    case (_, Nil) => -1
    case (a::_, b::_) if(a!=b) => a.compare(b)
    case (_::as, _::bs) => compare(as,bs)
  }

  override def equals(that: Any) = that match { 
    case Literal(po2, pr2) => po2==positive && pr2==predicate
    case _ => false
  }

  def negated = Literal(!positive, predicate)

  override def literals = TreeSet(this)
  override def variables = predicate.variables
}

//========================
//  Classical FOL Formula
//========================

case class Negation(formula: Formula) extends Formula { 
  override def toString: String = "¬"+formula.toString
  override def contains(term: Term) = formula.contains(term)
  override def signature = formula.signature
  override def literals = formula.literals
  override def variables = formula.variables
}

case class Disjunction(var disjuncts: Set[Formula]) extends Formula { 
  // Flatten possible nested disjunctions
  disjuncts = disjuncts.map(_ match { 
    case Disjunction(ds) => ds
    case x => Set(x)
  }).flatten

  override def toString: String = 
    if (disjuncts.isEmpty)
      "FALSE"
    else
      disjuncts.mkString("(", " || ", ")")

  override def contradiction = disjuncts.isEmpty
  override def contains(term: Term) = disjuncts.exists(_.contains(term))
  override def signature = disjuncts.flatMap(_.signature)
  override def literals = disjuncts.flatMap(_.literals)
  override def variables = disjuncts.flatMap(_.variables)

  override def equals(that: Any) = that match { 
    case Disjunction(ds) => disjuncts.sameElements(ds)
    case _ => false
  }

  override def hashCode = disjuncts.map(_.hashCode).sum
}

case class Conjunction(var conjuncts: Set[_ <: Formula]) extends Formula { 
  // Flatten possible nested conjunctions
  conjuncts = conjuncts.map(_ match { 
    case Conjunction(cs) => cs
    case x => Set(x)
  }).flatten
  
  override def toString: String = 
    if(conjuncts.isEmpty)
      "TRUE"
    else 
      conjuncts.mkString("(", " && ", ")")

  override def tautology = conjuncts.isEmpty
  override def contains(term: Term) = conjuncts.exists(_.contains(term))
  override def signature = conjuncts.flatMap(_.signature)
  override def literals = conjuncts.flatMap(_.literals)
  override def variables = conjuncts.flatMap(_.variables)

  override def equals(that: Any) = that match { 
    case Conjunction(cs) => conjuncts.sameElements(cs)
    case _ => false
  }

  override def hashCode = conjuncts.map(_.hashCode).sum
}

case class Implication(formula1: Formula, formula2: Formula) extends Formula { 
  override def toString: String = "("+formula1.toString+" -> "+formula2.toString+")"
  override def contains(term: Term) = formula1.contains(term) || formula2.contains(term)
  override def signature = formula1.signature ++ formula2.signature
  override def literals = formula1.literals ++ formula2.literals
  override def variables = formula1.variables ++ formula2.variables
}

case class Equivalence(formula1: Formula, formula2: Formula) extends Formula { 
  override def toString: String = "("+formula1.toString+" <-> "+formula2.toString+")"
  override def contains(term: Term) = formula1.contains(term) || formula2.contains(term)
  override def signature = formula1.signature ++ formula2.signature
  override def literals = formula1.literals ++ formula2.literals
  override def variables = formula1.variables ++ formula2.variables
}

case class ExistentialQuantification(variable: Variable, formula: Formula) extends Formula { 
  override def toString: String = "E"+variable.toString+"."+formula
  override def contains(term: Term) = variable==term || formula.contains(term)
  override def signature = formula.signature
  override def literals = formula.literals
  override def variables = formula.variables
}

case class UniversalQuantification(variable: Variable, formula: Formula) extends Formula{ 
  override def toString: String = "A"+variable.toString+"."+formula
  override def contains(term: Term) = variable==term || formula.contains(term)
  override def signature = formula.signature
  override def literals = formula.literals
  override def variables = formula.variables
}

//===========
//  Clauses
//===========
 
case class Clause(val _literals: Set[Literal]) extends Formula with Ordered[Clause] {

  val literalsSorted: SortedSet[Literal] = TreeSet.empty[Literal] ++ _literals

  override def literals: SortedSet[Literal] = { 
    assert(_literals!=null)
    assert(literalsSorted!=null)
    assert(literalsSorted.size==_literals.size, "Inconsistent ordering! "+_literals.toString+" vs. "+literalsSorted.toString)

    literalsSorted
  }

  // ensure symmetry of identity predicates
  // _literals = literals ++ literals.collect{ 
  //   case Literal(p, IdentityPredicate(x,y)) => new Literal(p, IdentityPredicate(y,x)) 
  // }

  def positivePart: Iterable[Literal] = literals filter{ _.positive }
  def negativePart: Iterable[Literal] = literals filter{ !_.positive }
  
  override def toString: String = literals.mkString(" || ")
  
  def contains(literal: Literal): Boolean = _literals.toSeq.contains(literal)
  def contains(term: Term): Boolean = literals.exists(_.predicate.args.exists(x => x==term || x.contains(x)))

  def isEmpty: Boolean = literals.isEmpty
  
  // returns a clause without the given literal
  def without(literal: Literal) = { 
    val result = literal match {
      case Literal(p, IdentityPredicate(x,y)) => new Clause(_literals.toSet[Literal] - literal - new Literal(p, IdentityPredicate(y,x)))
//      case _ =>   new Clause(_literals.toSet[Literal] - literal)
      case _ => new Clause(_literals.filterNot(l => l==literal))
    }

    assert(!result.contains(literal), literal.toString)

    result
  }
  
  def merge(other: Clause) = new Clause(literals ++ other.literals)
  
  def add(newLiterals: Iterable[Literal]) = new Clause(literals ++ newLiterals) 
  
  def terms = literals.flatMap(_.predicate.args.flatMap(_.subterms))

  // checks whether a tautology is directly visible (without unification)
  override def tautology = positivePart.exists(x => negativePart.exists(y => x.predicate==y.predicate))

  override def signature = literals.flatMap(_.signature)

  override def compare(that: Clause) = { 

    def inner(ls: Seq[Literal], ls2: Seq[Literal]): Int = (ls, ls2) match { 
      case (Nil, Nil) => 0
      case (Nil, set) => 1
      case (set, Nil) => -1
      case (set1, set2) => { 
	val result = set1.head.compare(set2.head)
	if(result==0)
	  inner(set1.tail, set2.tail)
	else
	  result
      }
    }

    inner(literals.toSeq, that.literals.toSeq)
  }

  override def variables = literals.toSeq.flatMap(_.variables).toSet
}


//==================
//  Formulas object
//==================

object FolHelpers { 
  def False = new Disjunction(Set())
  def True = new Conjunction(Set())

  def toFormula(clauses: Set[_ <: Clause]): Formula = { 
    simplify(new Conjunction(clauses.map(x => new Disjunction(x.literals.toSet[Formula]))))
  }

  def simplify(formula: Formula): Formula = formula match { 
    case Conjunction(cs) => { 
      val scs = cs.map(simplify)
      if(scs.exists(_.contradiction))
	FolHelpers.False
      else if(scs.size==1)
	scs.head
      else
	new Conjunction(scs)
    }
    case Disjunction(ds) => { 
      val sds = ds.map(simplify)
      if(sds.exists(_.tautology))
	FolHelpers.True
      else if(sds.size==1)
	sds.head
      else
	new Disjunction(sds)
    }
    case Negation(f1) => f1 match { 
      case Negation(f2) => simplify(f2)
      case _ => new Negation(simplify(f1))
    }
    case _ => formula
  }

  def differentVariables(p1: Predicate, p2: Predicate): Boolean = { 
    !p1.args.exists(x => p2.args.exists(y => x==y))
  }

  def implication(d: Disjunction) = { 
    val leftSideOption = d.disjuncts.collectFirst { 
      case l: Literal if !l.positive => l
      case n: Negation => n
    }
    leftSideOption match { 
      case Some(Literal(t, p)) => Implication(Literal(!t, p), disjunction(d.disjuncts - leftSideOption.get))
      case Some(Negation(p)) => Implication(p, disjunction(d.disjuncts - leftSideOption.get))
      case None => Implication(True, d)
    }
  }

  def disjunction(ds: Iterable[Formula]) = { 
    val set = ds.toSet[Formula]
    if(set.contains(True))
      True
    else { 
      def result = Disjunction(set) // automatically flattens
      if(result.disjuncts.size==1) 
	result.disjuncts.head
      else 
	result
    }
  }

  def conjunction(cs: Iterable[Formula]) = { 
    val set = cs.toSet[Formula]
    if(set.contains(False))
      False
    else { 
      def result = Conjunction(set) // automatically flattens
      if(result.conjuncts.size==1) 
	result.conjuncts.head
      else 
	result
    }
  }


  def negate(f: Formula): Formula = f match{ 
    case Literal(p, f) => Literal(!p, f)
    case Negation(f) => f
    case f => Negation(f)
  }


  def universal(v: Variable, f: Formula) = { 
    if(f.tautology)
      FolHelpers.True
    else if(f.contradiction)
      FolHelpers.False
    else
      new UniversalQuantification(v, f)
  }

  def existential(v: Variable, f: Formula) = { 
    if(f.tautology)
      FolHelpers.True
    else if(f.contradiction)
      FolHelpers.False
    else
      new ExistentialQuantification(v,f)
  }

}
