(* Written by Likai Liu; some minor modifications by Hongwei Xi *)

datatype BraunTree = E | B of BraunTree * BraunTree

(* make a braun tree with however number of elements *)

fun numBranchNode (E) = 0
  | numBranchNode (B(t1,t2)) = 1 + numBranchNode(t1) + numBranchNode(t2)

fun addOneNode (E) = B(E, E)
  | addOneNode (B(t1, t2)) =
    if numBranchNode(t1) > numBranchNode(t2) then B(t1, addOneNode(t2))
    else B(addOneNode(t1), t2)

fun makeBraunTreeFrom (t:BraunTree, 0) = t
  | makeBraunTreeFrom (t:BraunTree, left:int) =
    makeBraunTreeFrom (addOneNode(t), left - 1)

(* a tail-recursive implementation of the power function *)
fun pow (n:int, x:int): int =
    let
	fun aux (n, x, res) =
	    if n = 0 then res
	    else if n = 1 then x * res
		 else if n mod 2 = 0 then aux (n div 2, x * x, res)
		      else aux (n div 2, x * x, x * res)
    in
	aux (n, x, 1)
    end

fun listBraunTrees (height: int) : BraunTree list =
    let
	val base = pow(height - 1, 2)
	val templateTree = makeBraunTreeFrom(E, base)

	fun loop (start:int, stop:int, stack:BraunTree list,
		  lastTempTree:BraunTree) =
	    if (start > stop) then stack
	    else
		let
		    val newTree = addOneNode(lastTempTree)
		in
		    loop (start + 1, stop, stack @ [lastTempTree], newTree)
		end
    in
	loop (base, pow(height, 2) - 1, [], templateTree)
    end

fun printBraunTree (t: BraunTree): unit =
    let
	fun aux (E) = print "E"
	  | aux (B(t1, t2)) =
	    (print "B("; aux(t1); print ", "; aux(t2); print ")")
    in
	aux(t); print "\n"
    end