(* Binary-tree implementation of arbitrary-depth arrays *)
app load ["Word8"];

abstype atom = N  of Word8.word
             | FN of unit -> arr
             | FM of arr -> arr
             | FD of (arr * arr) -> arr
and     arr  = Lf of atom (* arrays as binary trees - SML doesn't allow arbitrarily nested lists *)
             | Nd of arr * arr (* (val, next) *)
             | Bx of arr * arr (* (shape, arr) *)
             | Zl (* Zilde *)
with
(* K exceptions *)
exception Domain;
exception Index;
exception BadStop;
exception NYI;
(* atom constructors *)
fun Word x = Lf(N(x));
fun Fun0 f = Lf(FN(f));
fun Fun1 f = Lf(FM(f));
fun Fun2 f = Lf(FD(f));
(* array constructor *)
fun Array []     = Zl
  | Array [x]    = x
  | Array (x::xs) = Nd(x,Array(xs)); (* TODO: automatically box? *)
(* return tree-element type *)
fun tt Zl        = "Zl"
  | tt (Lf(x))   = "Lf"
  | tt (Nd(x,y)) = "Nd"
  | tt (Bx(s,a)) = "Bx"
  | tt _         = raise Domain;
(* convert an array to a list of its leaves *)
fun leaves Zl        = []
  | leaves (Lf(x))   = [Lf(x)]
  | leaves (Nd(x,y)) = x::(leaves y)
  | leaves (Bx(s,a)) = leaves a;
(* unbox a boxed array *)
fun unbox (Bx(s,a)) = a
  | unbox (x)       = x;
(* return length of array *)
fun tally Zl          = 0w0
  | tally (Lf(a:atom)) = 0w1
  | tally (Nd(x,y))    = 0w1 + tally y
  | tally (Bx(s,a))    = tally (unbox a);
(* box an array *)
fun Box (a:arr) = Bx(Array([Word(tally a)]),a);
(* return shape of boxed array *)
(* TODO: do we only want this for boxed arrays?  do we want just tally for unboxed arrays? *)
(* TODO: Do we want a boxed array to be length 1 or the length of the contents? *)
fun shape (Bx(s,a)) = s
  | shape (Nd(x,y)) = Array([Word(tally(Nd(x,y)))]);
(* change shape of boxed array *)
fun reshape s (Bx(s,a)) = raise NYI
  | reshape s a = reshape s (Box(a));
(* print a string representation of an array *)
fun str a = 
  let fun st Zl    = "(Zilde)" (* TODO: box-drawing characters might be nice in the future *)
        | st (Lf(N(x)))  = StringCvt.padLeft #"0" 2 (Word8.fmt StringCvt.HEX x)
        | st (Lf(FN(x))) = "(Nilad)" (* TODO: would be nice to get actual function names *)
        | st (Lf(FM(x))) = "(Monad)"
        | st (Lf(FD(x))) = "(Dyad)"
        | st (Nd(x,y))   = (st x) ^ " " ^ (st y)
        | st (Bx(s,a))   = "[" ^ (st a) ^ "]"
  in (st a) ^ "\n"
  end;
fun p2 (Bx(s,a)) =
  let val l           = foldr (fn (x,y) => x*y) 0w1 (map (fn (Lf(N(x))) => x) (leaves a))
      fun d n []      = []
        | d 0 xs      = xs
        | d n (x::xs) = d (n-1) xs
      fun r _ []      = []
        | r 0 xs      = []
        | r n (x::xs) = x::(r (n-1) xs)
        (* FIXME: finish implementing *)
  in r 2 (leaves a)
  end;
(* return a string representation of the internal structure of an array *)
fun tree a =
  let fun r s n = concat (List.tabulate (n, fn _ => s));
  in let fun t (Zl)        n = (r " " n) ^ "ZILDE"
           | t (Lf(N(x)))  n = (r " " n) ^ "LEAF: " ^ StringCvt.padLeft #"0" 2 ((Word8.fmt StringCvt.HEX x))
           | t (Lf(FN(x))) n = (r " " n) ^ "LEAF: Nilad" (* TODO: would be nice to get actual function names *)
           | t (Lf(FM(x))) n = (r " " n) ^ "LEAF: Monad"
           | t (Lf(FD(x))) n = (r " " n) ^ "LEAF: Dyad"
           | t (Nd(x,y))   n = (r " " n) ^ "NODE:\n" ^ (r " " (n+1)) ^ (t x (n+1)) ^ ",\n" ^ (r " " (n+1)) ^ (t y (n+1))
           | t (Bx(x,y))   n = (r " " n) ^ "BOX:\n" ^ (r " " (n+1)) ^ (t y (n+1))
     in (t a 0) ^ "\n"
   end
  end;
(* call a nilad *)
(* TODO: higher-order functions *)
fun apply0 f = f(); (* TODO - what if it generates a whole array? *)
(* apply a monad to an array *)
(* TODO: higher-order functions *)
(* TODO: apply to an atom? *)
fun apply1 f Zl         = raise Domain (* TODO: is this right? *)
  | apply1 f (Lf(N(x))) = Lf(N(f x))
  | apply1 f (Nd(x,y))  = Nd((apply1 f x),(apply1 f y))
  | apply1 f (Bx(s,a))  = raise NYI; (* FIXME *)
(* apply a dyad to two arrays *)
(* TODO: higher-order functions *)
(* TODO: apply to atoms, or a mix of both?  do we need tagged functions for e.g. # that changes dependent on atom or array? *)
fun apply2 f x y = raise NYI; (* TODO *)
(* append an element to an array *)
fun append x a = Nd(x,a);
(* return an element at an index *)
fun index _ Zl          = Zl (* TODO: Is this the behaviour that we want?  Or should it raise an exception? *)
  | index 0w0 (Lf(x))   = Lf(x)
  | index _ (Lf(x))     = raise Index
  | index 0w0 (Nd(x,y)) = x
  | index i (Nd(x,y))   = index (i-0w1) y
  | index i (Bx(s,a))   = index i a;
(* take the first n elements from an array *)
fun take _ Zl          = Zl
  | take 0w0 a         = raise Index
  | take 0w1 (Lf(x))   = Lf(x)
  | take 0w1 (Bx(s,a)) = raise NYI (* TODO: implement taking from first axis like dyalog? *)
  | take 0w1 (Nd(x,y)) = x
  | take n (Nd(x,y))   = Nd(x,(take (n-0w1) y))
  | take n (Bx(x,y))   = raise NYI; (* TODO: see above *)
(* drop the first n elements from an array *)
fun drop _ Zl          = Zl
  | drop 0w0 a         = a
  | drop 0w1 (Nd(x,y)) = y
  | drop n (Nd(x,y))   = drop (n-0w1) y
  | drop _ (Bx(s,a))   = raise NYI; (* TODO: see notes on boxes on take above *)
(* reverse an array *)
fun reverse (Lf(x))   = Lf(x) (* FIXME: this doesn't behave properly, see str(take 0w4 (reverse (iota 0w8))); *)
  | reverse (Bx(s,a)) = Bx(s,a) (* TODO: boxing behaviour might be inconsistent across functions; check/correct asap before you get too deep into it *)
  | reverse (Nd(x,y)) = let fun r xs = case xs of [] => [] | (x::xs) => (r xs) @ [x]
                        in Array(r (leaves (Nd(x,y))))
                        end;
(* return the first n integers (zero-indexed) *)
fun iota n =
  let fun i 0w0 = raise Index (* TODO: is this the right exception? *)
        | i 0w1 = Word(0w0)
        | i n   = Nd(Word(n-0w1),(i (n-0w1)))
  in reverse (i n)
  end;
(* fun addD(x,Lf(STOP)) = x (* dyadic add *)
  | addD(Lf(N(x)),Lf(N(y))) = Lf(N(x+y))
  | addD(Nd(x,xs),Lf(N(y))) = Nd(Lf(N(x+y)),add(xs,Lf(N(y)))) (* FIXME: doesn't work because x might be a Nd - need to do cases for Nd(Lf(N(x)),xs) &c. *) 
  | addD(Lf(N(x)),Nd(y,ys)) = Nd(Lf(N(x+y)),add(Lf(N(x)),ys))
  | addD(Nd(x,xs),Nd(y,ys)) = raise NYI; (* TODO *) *)
end;

(* Test data *)
val a = Box(Array([Word(0w2),Word(0w3),Word(0w4),Word(0w5)]));
val b = Box(Array([Word(0w2),Box(Array([Word(0w3),Word(0w4)])),Word(0w5)]));