// TODO: remember to add your definition of trees from task 4

domain Tree {
    /* constructors */
    function Leaf(): Tree
    function Node(v: Int, lft: Tree, rgt: Tree): Tree

    function type(t: Tree): Int
    unique function type_Leaf(): Int
    unique function type_Node(): Int
    axiom type_of_Leaf {
    type(Leaf()) == type_Leaf()
    }
    axiom type_of_Node {
    forall v: Int, lft: Tree, rgt: Tree ::
    type(Node(v, lft, rgt)) == type_Node()
    }
    axiom all_types {
    forall t: Tree ::
    type(t) == type_Leaf() || type(t) == type_Node() 
    }

    /* deconstructors */
    
    function isLeaf(t: Tree): Bool
    function isNode(t: Tree): Bool
    function getVal(t: Tree): Int
    function getLeft(t: Tree): Tree
    function getRight(t: Tree): Tree

 
     /* correctness axioms */

    axiom isNode_correctness {
        forall t : Tree:: isNode(t) <==> type(t) == type_Node()  
    }

    axiom isLeaf_correctness {
        forall t : Tree:: isLeaf(t) <==> type(t) == type_Leaf() 
    }
    
    axiom getVal_correctness {
        forall val : Int, lft : Tree, rgt : Tree :: getVal(Node(val,lft,rgt)) == val 
    }

    axiom getLeft_correctness {
        forall val : Int, lft : Tree, rgt : Tree :: getLeft(Node(val,lft,rgt)) == lft  
    }

    axiom getRight_correctness {
       forall val : Int, lft : Tree, rgt : Tree :: getRight(Node(val,lft,rgt)) == rgt 
    }
}




function height(t: Tree): Int 
{
    (isLeaf(t)) 
    ? 0  
    : height(getLeft(t)) >= height(getRight(t))
    ? 1 + height(getLeft(t))
    : 1 + height(getRight(t))
}



method computeTreeHeight(t: Tree) returns (res: Int)
ensures res == height(t)
{

  
  if (isLeaf(t)){
    res := 0
  }else{
   
    var current: Seq[Tree]
    var next: Seq[Tree]
    current := Seq(t)
    next := Seq[Tree]()
    res := 0


    
    while (|current| > 0)
    invariant forall  x:Int :: (x >= 0 && x < |current|) ==> isNode(current[x]) //guarantees height relation between iterations
    invariant forall  x:Int :: ((x >= 0 && x < |current|) && (forall  y:Int :: (y >= 0 && y < |current|) ==> height(current[x]) >= height(current[y])))
              ==> height(t) == res + height(current[x]) //res == height(t) - (height of biggest subtree)
    invariant (|current| == 0) ==> height(t) == res //needed for the final stretch

    {
      res := res + 1
      next := Seq[Tree]()

      var old_current : Seq[Tree] := current //for invariants and asserts only

      while (|current| > 0)
      invariant forall  x:Int :: (x >= 0 && x < |current|) ==> isNode(current[x])
      invariant (|current| == 0 && |next| == 0 ) ==> height(t) == res 
      
      /*
      all trees in old_current have height 1, (case when all nodes in current were childless so we end up with |current| = |next| = 0)
      or the tree in old_current of maximal height is still in current,
      or the children, hence also child tree of maximal height, of the tree of maximal height from old_current was added to next
      */
      invariant forall  x:Int :: ((x >= 0 && x < |old_current|) && height(t) == res + height(old_current[x]) - 1) ==> (
        (forall  k:Int :: ((k >= 0 && k < |old_current|) ==> !isNode(getLeft(old_current[k])) && !isNode(getRight(old_current[k])))) ||
        (|current| > 0 && (forall  z:Int :: ((z >= 0 && z < |current|) && forall  y:Int :: (y >= 0 && y < |current|) ==> height(current[z]) >= height(current[y]))
              ==> height(t) == res + height(current[z])-1))
              ||
        (|next| > 0 && (forall  v:Int :: ((v >= 0 && v < |next|) && forall  y:Int :: (y >= 0 && y < |next|) ==> height(next[v]) >= height(next[y]))
           ==> height(t) == res + height(next[v])))
      )

      //if all trees in old_current have height 1, the same is true for current
      invariant (forall  k:Int ::  ((k >= 0 && k < |old_current|) ==> !isNode(getLeft(old_current[k])) && !isNode(getRight(old_current[k])))) ==> 
                (forall  k:Int ::  ((k >= 0 && k < |current|) ==> !isNode(getLeft(current[k])) && !isNode(getRight(current[k])))) 

      //if none of the nodes in old_current are childless, nothing is added to next 
      invariant (forall  k:Int :: ((k >= 0 && k < |old_current|) ==> !isNode(getLeft(old_current[k])) && !isNode(getRight(old_current[k])))) ==> |next| == 0


      {

        var node : Tree := current[0] 

        POP(current)

        assert height(t) ==  res + height(node)-1 ==> (height(t) == res + height(getLeft(node)) || height(t) == res + height(getRight(node))) 

        if (isNode(getLeft(node))){
          assert false							// <-- HERE 
          PUSH(next, getLeft(node))
        }
        if (isNode(getRight(node))){
          PUSH(next, getRight(node))
        }

       
      }

      current := next

      
    }

  }

}

define PUSH(stck, v) {
  stck := Seq(v) ++ stck
}



define POP(stck) {
  stck := stck[1..]
}