include "pset06_definitions.dfy"

module Pset6 {
  import opened Defs

  // In this pset, you will work with various functions around Boolean formulas that requires explicit 'decreases' clauses.

  // First, let's start with SATAux from the previous pset. Note that we used a fuel parameter so that Dafny is able to infer that fuel decreases by 1 in each recursive call. This is redundant, and you are asked to supply the decreasing arguments in DecreaseForSATAux2 to get rid of this fuel parameter.
  function SATAux(fuel: nat, e: BooleanFormula, assignment: State): Option<State>
    requires |BFDomain(e)| == fuel
    ensures SATAux(|BFDomain(e)|, e, assignment) == SATAux2(e, assignment)
  {
    if fuel == 0 then
      if EvalBF(e, assignment) then Some(assignment) else None
    else
      match SATAux(fuel - 1, SubstBF(e, GetVariable(e), true), assignment[GetVariable(e) := true])
      case Some(a) => Some(a)
      case None => SATAux(fuel - 1, SubstBF(e, GetVariable(e), false), assignment[GetVariable(e) := false])
  }

  function DecreaseForSATAux2(e: BooleanFormula, assignment: State): PlaceHolderType
  {
    // Your Code
    PlaceHolderValue
  }

  function SATAux2(e: BooleanFormula, assignment: State): (opt: Option<State>)
    ensures opt.Some? ==> BFDomain(e) <= opt.value.Keys
    ensures opt.Some? ==> assignment.Keys <= opt.value.Keys
    ensures opt.Some? ==> forall k :: k in assignment.Keys ==> k !in BFDomain(e) ==> assignment[k] == opt.value[k]
    decreases DecreaseForSATAux2(e, assignment)
  {
    if BFDomain(e) == {} then
      if EvalBF(e, assignment) then Some(assignment) else None
    else
      match SATAux2(SubstBF(e, GetVariable(e), true), assignment[GetVariable(e) := true])
      case Some(a) => Some(a)
      case None => SATAux2(SubstBF(e, GetVariable(e), false), assignment[GetVariable(e) := false])
  }

  // The enumeration method from the last pset to solve SAT problems is probably the most naive and the most expensive approach that one can come up with. In practice, general Boolean formulas are normalized into their equivalent CNF (Conjunction Normal Form) versions, and we use algorithms with heuristics on the CNF representation to speed up solving the common cases. You are asked to first work through a few steps that are involved in CNF conversion, and we will look at the basic enumeration method on CNF.

  // The steps to translate from a general Boolean formula are 1) convert to NNF (Negation Normal Form), 2) linearize conjunction and disjunction, and 3) distribute. The implementations of all the intermediate steps are provided, and you are asked to fill in the 'decrease' clauses for some of them. Finally, you need to put these intermediate implementations together to complete the final algorithm CNFConversion.

  // NNF conversion takes a general Boolean formula and pushes all the negations to innermost layers by De Morgan's laws and double negation. Dafny cannot automatically infer the 'decreases' clause for this implementation. Fill in DecreaseForNNFConversion so that Dafny can check termination.
  // Example of NNF: !( x && (!y)) -->  (!x || y)
  function NNFConversion(e: BooleanFormula): BooleanFormula
    ensures IsNNF(NNFConversion(e))
    ensures BFDomain(NNFConversion(e)) == BFDomain(e)
    ensures forall st: State :: BFDomain(e) <= st.Keys ==> EvalBF(NNFConversion(e), st) == EvalBF(e, st)
    decreases DecreaseForNNFConversion(e)
  {
    match e
    case And(a, b) => And(NNFConversion(a), NNFConversion(b))
    case Or(a, b) => Or(NNFConversion(a), NNFConversion(b))
    case Not(Not(a)) => NNFConversion(a)
    case Not(And(a, b)) => Or(NNFConversion(Not(a)), NNFConversion(Not(b)))
    case Not(Or(a, b)) => And(NNFConversion(Not(a)), NNFConversion(Not(b)))
    case _ => e
  }

  function DecreaseForNNFConversion(e: BooleanFormula): PlaceHolderType
  {
    // Your Code
    PlaceHolderValue
  }

  // Linearization takes a Boolean formula in NNF and linearizes nested conjunctions and disjunctions (but no distribution yet). This normalization makes it easier to deal with in later steps. Fill in DecreaseForLinearize so that Dafny can check termination. You may find it helpful to add new ensures clauses for Linearize.
  // Example of Linearization: (x && (y && z) && ((u && v) && w)) --> (x && (y && (z && (u && (v && w))))
  function Linearize(e: BooleanFormula): BooleanFormula
    requires IsNNF(e)
    ensures IsNNF(Linearize(e))
    ensures IsLinear(Linearize(e))
    ensures BFDomain(e) == BFDomain(Linearize(e))
    ensures forall st: State :: BFDomain(e) <= st.Keys ==> EvalBF(Linearize(e), st) == EvalBF(e, st)
    decreases DecreaseForLinearize(e)
  {
    match e
    case And(a, b) =>
      var a_linearized := Linearize(a);
      match a_linearized {
        case And(a0, b0) =>
          And(a0, Linearize(And(b0, b)))
        case _ =>
          And(a_linearized, Linearize(b))
      }
    case Or(a, b) =>
      var a_linearized := Linearize(a);
      match a_linearized {
        case Or(a0, b0) =>
          Or(a0, Linearize(Or(b0, b)))
        case _ =>
          Or(a_linearized, Linearize(b))
      }
    case _ => e
  }

  function DecreaseForLinearize(e: BooleanFormula): PlaceHolderType
  {
    // Your Code
    PlaceHolderValue
  }

  // Helper Lemma for Distribution
  lemma LinearCNFImpDisjunction(a: BooleanFormula, b: BooleanFormula)
    requires IsCNF(And(a, b))
    requires IsLinear(And(a, b))
    ensures FreeOfConjunction(a)
    ensures IsLinear(a)
  {
  }

  // Last, we need to push all disjunction to the leaves of the expression. This is done by the distribution law (i.e. (x && y) || z --> (x || z) && (y || z)). This part is implemented for you, but you will need it to complete CNFConversion.
  function DistributeOverDisjunctionAux(a: BooleanFormula, b: BooleanFormula): BooleanFormula
    requires FreeOfConjunction(a)
    requires IsNNF(a)
    requires IsLinear(b)
    requires IsCNF(b)
    ensures IsCNF(DistributeOverDisjunctionAux(a, b))
    ensures BFDomain(DistributeOverDisjunctionAux(a, b)) == BFDomain(a) + BFDomain(b)
    ensures forall st: State :: BFDomain(a) <= st.Keys ==> BFDomain(b) <= st.Keys ==> EvalBF(DistributeOverDisjunctionAux(a, b), st) == (EvalBF(a, st) || EvalBF(b, st))
  {
    match b
    case Const(x) => Or(a, b)
    case Var(name) => Or(a, b)
    case And(c, d) =>
      LinearCNFImpDisjunction(c, d);
      assert IsCNF(Or(a, c));
      assert IsCNF(DistributeOverDisjunctionAux(a, d));
      And(Or(a, c), DistributeOverDisjunctionAux(a, d))
    case Or(c, d) => Or(a, b)
    case Not(c) => assert FreeOfConjunction(c); Or(a, b)
  }

  function DistributeOverDisjunction(a: BooleanFormula, b: BooleanFormula): BooleanFormula
    requires IsLinear(a)
    requires IsLinear(b)
    requires IsCNF(a)
    requires IsCNF(b)
    ensures IsCNF(DistributeOverDisjunction(a, b))
    ensures BFDomain(DistributeOverDisjunction(a, b)) == BFDomain(a) + BFDomain(b)
    ensures forall st: State :: BFDomain(a) <= st.Keys ==> BFDomain(b) <= st.Keys ==> EvalBF(DistributeOverDisjunction(a, b), st) == (EvalBF(a, st) || EvalBF(b, st))
  {
    match a
    case And(c, d) =>
      LinearCNFImpDisjunction(c, d);
      assert FreeOfConjunction(c);
      And(DistributeOverDisjunctionAux(c, b), DistributeOverDisjunction(d, b))
    case _ => DistributeOverDisjunctionAux(a, b)
  }

  // Last, fill in the CNF conversion so all the postconditions hold.
  function CNFConversion(e: BooleanFormula): BooleanFormula
    requires IsNNF(e)
    ensures IsCNF(CNFConversion(e))
    ensures IsLinear(CNFConversion(e))
    ensures BFDomain(CNFConversion(e)) == BFDomain(e)
    ensures forall st: State :: BFDomain(e) <= st.Keys ==> EvalBF(CNFConversion(e), st) == (EvalBF(e, st))
  {
    // Your Code
    Var("PlaceHolder")
  }

  /*
    Now we have completed all the steps to translate any Boolean formula to its CNF equivalent (and also proved its correctness implicitly by the ensure clauses!). We will now switch gears and work with a more structured modeling of any CNF formula instead of type `BooleanFormula`.
    Type `CNF` is a 2-dimensional list of Literals, where the inner dimension represents disjunctions of literals (called clauses), and the outer dimension represents the conjunctions of clauses. Literal construction takes a sign and a string, where the sign represents its polarity and the string is just its name.
    For example, !x is constructed by Lit(false, "x"), while y is constructed by Lit(true, "y"). In addition, (x || y || !z) && (!x || u) && (z) has the following representation:
    Cons(
      Cons( Lit(true, "x"), Cons( Lit(true, "y"), Cons(Lit(false, "z"), Nil ) ) )
      Cons( Lit(false, "x"), Cons( Lit(true, "u"), Nil ) )
      Cons( Lit(true, "z"), Nil )
      Nil
    )
  */

  // Let's start by building the essential function that returns the domain of a CNF. Fill in DecreaseForCNFDomain
  function CNFDomain(c: CNF): set<string>
    decreases DecreaseForCNFDomain(c)
  {
    match c
    case Nil => {}
    case Cons(Nil, ls) => CNFDomain(ls)
    case Cons(Cons(Lit(_, name), ls'), ls) => {name} + CNFDomain(Cons(ls', ls))
  }

  function DecreaseForCNFDomain(c: CNF): PlaceHolderType
  {
    // Your Code
    PlaceHolderValue
  }

  // The Evaluation of a CNF says, if the CNF has no clause, then its value is true. If it has an empty clause, which means none of the literals that were contained in this clause was evaluated to true under the input assignments, hence this clause is not satisfiable, then the value is false. Otherwise, the CNF is inductively evaluated literal-by-literal until one of the above two base cases.
  // Fill in DecreaseForEvalCNF. You will likely want to use HelperLemmaForEvalCNF to inject some useful information that Dafny can use to construct this function.
  function EvalCNF(c: CNF, assignment: map<string, bool>): bool
    requires CNFDomain(c) <= assignment.Keys
    decreases DecreaseForEvalCNF(c, assignment)
  {
    HelperLemmaForEvalCNF(c, assignment);
    match c
    case Nil => true
    case Cons(Nil, cnf) => false
    case Cons(Cons(Lit(sign, name), clause), cnf) =>
      if sign == assignment[name] then
      EvalCNF(cnf, assignment) else EvalCNF(Cons(clause, cnf), assignment)
  }

  function DecreaseForEvalCNF(c: CNF, assignment: map<string, bool>): PlaceHolderType
  {
    PlaceHolderValue
  }

  lemma HelperLemmaForEvalCNF(c: CNF, assignment: map<string, bool>)
  {
  }

  // The Substitution of a CNF changes the input CNF based on the variable and the value it is replaced with. Note that there are no Boolean constants in our model for CNF, since it is straightforward to propagate a constant as shown below.
  // Fill in DecreaseForSubstCNF. You will likely want to use HelperLemmaForSubstCNF to inject some useful information.
  function SubstCNF(c: CNF, variable: string, value: bool): CNF
    ensures variable !in CNFDomain(SubstCNF(c, variable, value))
    ensures CNFDomain(c) + {variable} > CNFDomain(SubstCNF(c, variable, value))
    decreases DecreaseForSubstCNF(c, variable, value)
  {
    HelperLemmaForSubstCNF(c, variable, value);
    match c
    case Nil => Nil
    case Cons(Nil, cnf) => Cons(Nil, SubstCNF(cnf, variable, value))
    case Cons(Cons(Lit(sign, name), clause), cnf) =>
      if variable == name then
      (
        if sign == value then
        SubstCNF(cnf, variable, value) else SubstCNF(Cons(clause, cnf), variable, value)
      )
      else
      (
        match SubstCNF(Cons(clause, cnf), variable, value)
        case Nil => Cons(Cons(Lit(sign, name), Nil), Nil)
        case Cons(clause, cnf) => Cons(Cons(Lit(sign, name), clause), cnf)
      )
  }

  function DecreaseForSubstCNF(c: CNF, variable: string, value: bool): PlaceHolderType
  {
    PlaceHolderValue
  }

  lemma HelperLemmaForSubstCNF(c: CNF, variable: string, value: bool)
  {
  }

  function FirstVariable(c: CNF): string
    requires CNFDomain(c) != {}
    ensures FirstVariable(c) in CNFDomain(c)
  {
    var Cons(h, tl) := c;
    match h
    case Nil => FirstVariable(tl)
    case Cons(Lit(sign, name), _) => name
  }

  // Last, fill in DecreaseForCNFSATAux as the decreases clause for CNFSATAux.
  function CNFSATAux(c: CNF, assignment: map<string, bool>): (opt: Option<map<string, bool>>)
    ensures opt.Some? ==> assignment.Keys <= opt.value.Keys
    decreases DecreaseForCNFSATAux(c, assignment)
  {
    if CNFDomain(c) == {} then
      match c
      case Nil => Some(assignment)
      case Cons(Nil, _) => None
    else
      match CNFSATAux(SubstCNF(c, FirstVariable(c), true), assignment[FirstVariable(c) := true])
      case Some(a) => Some(a)
      case None => CNFSATAux(SubstCNF(c, FirstVariable(c), false), assignment[FirstVariable(c) := false])
  }

  function DecreaseForCNFSATAux(c: CNF, assignment: map<string, bool>): PlaceHolderType
  {
    PlaceHolderValue
  }

}