// F# algebraic expression trees, with variables, factoring and re-balance

open System;
open System.Collections.Generic;

// This time there are variable expressions in a expr tree, so we can
// represent "x+2".  Variables can only be evaluated given an enviornment,
// which are bindings of the variables to values.  We will use a hashmap,
// or System.Collections.Generic.Dictionary<string,int> to represent the
// environment.  This representation is OK for global variables.  However,
// if we wanted to implement proper scoping, it would be inadequate.  We
// would need a stack combined with a representation of CLOSURES that
// can capture the stack in a specific state.

type expr = Num of int | Plus of expr*expr | Times of expr*expr | Var of string;;

let Global = Dictionary<string,int>();;  // Global variable bindings only
Global.["x"] <- 4;   // create a sample binding
Global.["y"] <- 7;
Global.["z"] <- 0;

// evaluate expression e under environment env:
let rec eval(e,env:Dictionary<string,int>) = 
  match e with
   | Num(x) -> x
   | Plus(x,y) -> eval(x,env) + eval(y,env)
   | Times(x,y) -> eval(x,env) * eval(y,env);
   | Var(x) -> env.[x];;

let rec tostring = function  // print formula in infix notation:
  | Num(x) -> string(x)
  | Plus(Plus(a,b),Plus(c,d)) ->  // drop extra parentheses (shallow)
      "("+tostring(a)+" + "+tostring(b)+" + "+tostring(c)+" + "+tostring(d)+")"
  | Plus(Plus(x,y),z) -> 
      "("+tostring(x)+" + "+tostring(y)+" + "+tostring(z)+")"
  | Plus(x,Plus(y,z)) ->
      "("+tostring(x)+" + "+tostring(y)+" + "+tostring(z)+")"
  | Plus(x,y) ->  "("+ tostring(x) + " + " + tostring(y) + ")"
  | Times(x,y) -> tostring(x) + tostring(y) // prints x*y as xy
  | Var(x) -> x;;

// sample tree for  (2+x) * y:
let t = Times(Plus(Num(2),Var("x")),Var("y"))
let t2 = Times(Plus(Num(1+1),Var "x" ),Var "y")
Console.WriteLine((t=t2));  // checking that = works on trees

Console.WriteLine(eval(t,Global));
Console.WriteLine(tostring(t));;

// Factor a formula into a polynomial, or additive normal form:
// (x+2)(y+x) = (x+2)y + (x+2)x = xy+2y+xx+2x, want this form

// this is tricky to write because you could have something like
// a(b+c)*(d+e)b, where the addition is deep inside the multiplications.
// you also must avoid infinite recursive calls forever, when something
// is not changed, like in ab.
let rec anf e = 
  match e with
   | Plus(x,y) -> Plus(anf(x),anf(y))
   | Times(x,Plus(y,z)) -> anf(Plus(Times(x,y),Times(x,z)))
   | Times(Plus(x,y),z) -> anf(Plus(Times(x,z),Times(y,z)))
   | Times(x,y) -> 
      let ax,ay = anf(x),anf(y)
      if ax=x && ay=y then Times(x,y) else anf(Times(ax,ay)) 
   | x -> x;; // default doesn't change e

let t3 = Times(Times(Var "a",Plus(Var("b"),Var "c")),Times(Plus(Var "d",Var "e"),Var("b")));

Console.WriteLine(tostring(anf(t)))  // (2+x)y was original t
Console.WriteLine(tostring(t3))
Console.WriteLine(tostring(anf(t3)))

// Now that we can rewrite a formula in additive normal form, we can also
// rebalance the tree: a+(b+(c+d)) should become (a+b) + (c+d)

// calculating the depth of height of tree, if we keep the tree
// balanced, then depth will never need more than O(log n) nested
// recursive calls:
let rec depth = function
  | Num(_) | Var(_) -> 1
  | Plus(x,y) | Times(x,y) -> Math.Max(depth(x),depth(y))+1;;

let rotateLeft = function
  | Plus(a,Plus(b,c)) -> Plus(Plus(a,b),c)
  | x -> x;;
let rotateRight = function
  | Plus(Plus(b,c),a) -> Plus(b,Plus(c,a))
  | x -> x;;

let rec balance(tree) =
  match tree with
    | Plus(a,b) ->
        let da,db = depth(a),depth(b)
        if da+1<db then balance(rotateLeft(tree))  // yikes! nested if-else
        else if db+1<da then balance(rotateRight(tree))
        else tree
    | x -> x

let t4 = Plus(Var "a",Plus(Var "b",Plus(Var "c",Plus(Var "d",Plus(Var "e",Var "f")))));
Console.WriteLine(tostring(t4))
Console.WriteLine(depth(t4))
let t5 = balance(t4)
Console.WriteLine(tostring(t5))
Console.WriteLine(depth(t5))