// Fill in a good loop invariant for our workhorse `IsolateEqns`.
// If you have time to complete the proof in terms of this invariant,
// more power to you!  However, we'll probably have to leave that task to
// the code we'll share for you to check out after class.
ghost predicate IsolateEqnsInvariant(eqs: Eqns, isolated: seq<(Var, Exp)>, not_isolated: set<Eqn>) {
  // Any variable on the LHS of an isolated equation should not appear in not_isolated.
  // Any variable on the LHS of an isolated equation should not appear in its own equation or anywhere in any *later* isolated equation.
  // An assignment satisfies eqns iff it satisfies isolated + not_isolated.
}


// Here is the implementation of the solver, which you shouldn't need to consult
// too carefully to write your invariant.  In fact, feel free to describe properties
// as English in comments, since it would take too long to get comfortable with all
// of the following code.

datatype Option<T> = None | Some(value: T)

// Here is the syntax of systems of equations.  The only unusual part is that we
// build in, to syntax nodes for variables, multiplication by a constant.  Also
// note that we use Dafny's real-number type, which for our purposes is more
// accurately considered to be a type of rational numbers, where the operations
// we use can all be computed exactly.
type Var = string
datatype Exp = Const(n: real) | Var(n: real, x: Var) | Plus(e1: Exp, e2: Exp)
type Eqn = (Exp, Exp)
type Eqns = set<Eqn>

// Our goal in solving a system of equations is to find a variable assignment
// that satisfies it.  A more general solver would find a compact description of
// *all* such assignments rather than just one, but we kept the goal more modest
// to simplify verification.
type Asgn = Var -> real

// Then evaluation of different kinds of syntax w.r.t. assignments is
// straightforward.

function EvalExp(asgn: Asgn, e: Exp): real {
  match e
  case Const(n) => n
  case Var(n, x) => n * asgn(x)
  case Plus(e1, e2) => EvalExp(asgn, e1) + EvalExp(asgn, e2)
}

ghost predicate EvalEqn(asgn: Asgn, eq: Eqn) {
  EvalExp(asgn, eq.0) == EvalExp(asgn, eq.1)
}

ghost predicate EvalEqns(asgn: Asgn, eqs: Eqns) {
  forall eq :: eq in eqs ==> EvalEqn(asgn, eq)
}

// We also want to explain which variables appear in a piece of syntax.

function ExpVars(e: Exp): set<Var> {
  match e
  case Const(_) => {}
  case Var(_, x) => {x}
  case Plus(e1, e2) => ExpVars(e1) + ExpVars(e2)
}

function EqnVars(eq: Eqn): set<Var> {
  ExpVars(eq.0) + ExpVars(eq.1)
}

// Now some key helper functions over syntax, used in our solver.
// We build specifications into these definitions and don't bother to narrate
// the details, since the overall proof is quite involved (and taxes Dafny's
// ability to find the proofs without timing out!).
// First, we multiply by a constant, which manifests as multiplying similarly in
// every spot where a constant appears.

function Multiply(k: real, e: Exp): Exp
  ensures ExpVars(Multiply(k, e)) == ExpVars(e)
  ensures forall asgn :: EvalExp(asgn, Multiply(k, e)) == k * EvalExp(asgn, e)
{
  match e
  case Const(n) => Const(k * n)
  case Var(n, x) => Var(k * n, x)
  case Plus(e1, e2) => Plus(Multiply(k, e1), Multiply(k, e2))
}

// Next, we need to be able to substitute equations of the form "x = e", for "e"
// an expression not mentioning "x".  We need different operations for different
// kinds of syntax to substitute.

// This next variant operates in a system of equations, building in the removal
// of one equation before substituting.  (In intended usage, that equation will
// be one that we manipulated into the form "x = e" before substituting.)
function SubstituteEqns(eqs: Eqns, eq: Eqn, x: Var, e: Exp): Eqns {
  set eq' | eq' in eqs - {eq} :: SubstituteEqn(x, e, eq')
}

ghost function OverrideAsgn(asgn: Asgn, x: Var, v: real): Asgn {
  y => if y == x then v else asgn(y)
}

function SubstituteEqn(x: Var, e: Exp, eq: Eqn): Eqn
  ensures EqnVars(SubstituteEqn(x, e, eq)) <= EqnVars(eq) - {x} + ExpVars(e)
  ensures forall asgn :: EvalEqn(asgn, SubstituteEqn(x, e, eq)) == EvalEqn(OverrideAsgn(asgn, x, EvalExp(asgn, e)), eq)
{
  (SubstituteExp(x, e, eq.0), SubstituteExp(x, e, eq.1))
}

function SubstituteExp(x: Var, e1: Exp, e2: Exp): Exp
  ensures ExpVars(SubstituteExp(x, e1, e2)) <= ExpVars(e2) - {x} + ExpVars(e1)
  ensures forall asgn :: EvalExp(asgn, SubstituteExp(x, e1, e2)) == EvalExp(OverrideAsgn(asgn, x, EvalExp(asgn, e1)), e2)
{
  match e2
  case Const(_) => e2
  case Var(k, y) => if y == x then Multiply(k, e1) else e2
  case Plus(e21, e22) => Plus(SubstituteExp(x, e1, e21), SubstituteExp(x, e1, e22))
}

// Note that this style of substitution leaves rather large expressions, because
// it doesn't simplify by combining like terms.  This shortcut is another
// compromise to simplify verification.  Adding simplification is left as an
// exercise to the reader!  It certainly improves performance by a lot.

// Here is our overall solver.  It first *isolates* by rewriting each equation
// into some "x = e", where "x" does not appear in "e".  From that point, it is
// relatively easy to choose values, by evaluating the rewritten equations in
// reverse order.  Note that this method returns an Option type, allowing it to
// give up and return None in any situation.  We tried to write the code to
// return Some whenever there exists a solution to the equations, though we know
// one place where we cut a corner and that isn't true.  (Can you spot it?)
method SolveEqns(eqs: Eqns) returns (r: Option<Asgn>)
  ensures r.Some? ==> EvalEqns(r.value, eqs)
{
  var isolated := IsolateEqns(eqs);
  if isolated.None? {
    return None;
  }
  var asgn := ChooseValues(isolated.value, eqs);
  assume EvalEqns(asgn, eqs);
  return Some(asgn);
}

// The code for this one is less important, so we don't include it here.
method ChooseValues(isolated: seq<(Var, Exp)>, ghost eqs: Eqns) returns (r: Asgn)

// Here is the first main helper method.  We'll present its key loop invariant
// later.
method IsolateEqns(eqs: Eqns) returns (r: Option<seq<(Var, Exp)>>)
  // A spec is missing here.
{
  var isolated : seq<(Var, Exp)> := [];
  // These are the equations already rewritten as "x = e".
  var not_isolated := eqs;
  // These are the equations we didn't yet rewrite.

  while |not_isolated| > 0
    invariant IsolateEqnsInvariant(eqs, isolated, not_isolated)
  {
    var eq :| eq in not_isolated;
    if x :| x in EqnVars(eq) {
      // There is a variable appearing in the current equation.
      var eo := IsolateVarEqn(x, eq);
      match eo
      case None => return None;
      case Some(e) =>
        isolated, not_isolated := IsolateEqnsBody(eqs, isolated, not_isolated, eq, x, e);
    } else {
      // Weird: an equation with no variables.  It's most likely contradictory.
      return None;
    }
    assume false; // Included here to keep this file self-sufficient.
  }

  return Some(isolated);
}

// A helper to that helper: isolate a given variable within an equation.
function IsolateVarEqn(x: Var, eq: Eqn): Option<Exp>
  ensures match IsolateVarEqn(x, eq)
    case None => true
    case Some(e) => (ExpVars(e) == EqnVars(eq) - {x} &&
      forall asgn : Asgn :: EvalEqn(asgn, eq) == (asgn(x) == EvalExp(asgn, e)))
{
  var (n1, e1) := IsolateVarExp(x, eq.0);
  var (n2, e2) := IsolateVarExp(x, eq.1);
  var n := n1 - n2;
  if n == 0.0 then
    None
    // This is one fundanental place where the solver gives up on finding a
    // solution!  To proceed, we would need to divide by zero.
  else
    var e := Plus(Multiply(1.0/n, e2), Multiply(-1.0/n, e1));
    IsolateVarEqn'(eq, x, n1, e1, n2, e2);
      
    Some(e)
}

// This helper is for isolating a variable out of an expression (one side of an
// equation).
function IsolateVarExp(x: Var, e: Exp): (real, Exp)
  ensures var (n, e') := IsolateVarExp(x, e);
    ExpVars(e') == ExpVars(e) - {x} &&
    forall asgn :: EvalExp(asgn, e) == n * asgn(x) + EvalExp(asgn, e')
{
  match e
  case Const(_) => (0.0, e)
  case Var(n, y) => if y == x then (n, Const(0.0)) else (0.0, e)
  case Plus(e1, e2) =>
    var (n1, e1') := IsolateVarExp(x, e1);
    var (n2, e2') := IsolateVarExp(x, e2);
    (n1 + n2, Plus(e1', e2'))
}

// We have to do more explicit lemma work than we're used to in Dafny, because
// ironically, though we are verifying a *linear* equation solver, the fact we
// need to prove here uses *nonlinear* arithmetic, which Dafny struggles to
// handle efficiently and predictably.
lemma IsolateVarEquation(x: real, e1: real, n1: real, e1': real, e2: real, n2: real, e2': real)
  requires e1 == n1 * x + e1'
  requires e2 == n2 * x + e2'
  requires n1 - n2 != 0.0
  ensures (e1 == e2) == (x == (1.0 / (n1 - n2)) * e2' + (-1.0 / (n1 - n2)) * e1')
{
}

lemma IsolateVarEqn'(eq: Eqn, x: Var, n1: real, e1: Exp, n2: real, e2: Exp)
  requires forall asgn :: EvalExp(asgn, eq.0) == n1 * asgn(x) + EvalExp(asgn, e1)
  requires forall asgn :: EvalExp(asgn, eq.1) == n2 * asgn(x) + EvalExp(asgn, e2)
  requires n1 - n2 != 0.0
  ensures forall asgn: Asgn :: EvalEqn(asgn, eq) == (asgn(x) == EvalExp(asgn, Plus(Multiply(1.0/(n1-n2), e2), Multiply(-1.0/(n1-n2), e1))))
{
  forall asgn: Asgn ensures EvalEqn(asgn, eq) == (asgn(x) == EvalExp(asgn, Plus(Multiply(1.0/(n1-n2), e2), Multiply(-1.0/(n1-n2), e1)))) {
    IsolateVarEquation(asgn(x), EvalExp(asgn, eq.0), n1, EvalExp(asgn, e1),
      EvalExp(asgn, eq.1), n2, EvalExp(asgn, e2));
  }
}

// We pushed most of the action of the loop body into this method.
method IsolateEqnsBody(eqs: Eqns, old_isolated: seq<(Var, Exp)>, old_not_isolated: set<Eqn>,
  eq: Eqn, x: Var, e: Exp)
  returns (isolated: seq<(Var, Exp)>, not_isolated: set<Eqn>)
  requires eq in old_not_isolated
  requires x in EqnVars(eq)
  requires ExpVars(e) == EqnVars(eq) - {x}
  requires forall asgn : Asgn :: EvalEqn(asgn, eq) == (asgn(x) == EvalExp(asgn, e))
  requires IsolateEqnsInvariant(eqs, old_isolated, old_not_isolated)
{
  isolated := old_isolated + [(x, e)];
  not_isolated := SubstituteEqns(old_not_isolated, eq, x, e);
}