(*		Generators in OCaml

As the first attempt, we use the general theory: 
   yield = exceptions + non-determinism 

which we will optimize later. For non-determinism, we use HANSEI. For
exceptions -- native OCaml exceptions.

The general theory applies immediately: It just works.
*)

(* When working at the top-level, do:
#load "prob.cma";;
*)

open ProbM;; (* The HANSEI library *)
open Ptypes;;

(* Define yield, as derived in the general theory *)
let yield x = if flip 0.5 then raise x
;;

(* That is all needed to define generators *)

(* ------------------------------------------------------------------------ *)
(* Examples: in-order and post-order traversals of a tree *)

(* A few preliminaries: define the tree and build a sample one *)

type label = int
type tree = Leaf | Node of label * tree * tree;;

let make_full_tree depth =
  let rec loop label = function
    | 0 -> Leaf
    | n -> Node (label, loop (2*label) (pred n), loop (2*label+1) (pred n))
  in loop 1 depth
;;
(* val make_full_tree : int -> tree = <fun> *)

let tree1 = make_full_tree 3;;

(*
val tree1 : tree =
  Node (1, Node (2, Node (4, Leaf, Leaf), Node (5, Leaf, Leaf)),
           Node (3, Node (6, Leaf, Leaf), Node (7, Leaf, Leaf)))
*)


(* Declare the exception to carry the yielded value *)
exception YT of int;;

(* The traversal code itself *)
let rec in_order2 = function
  | Leaf -> ()
  | Node (label, left, right) ->
      in_order2 left;
      yield (YT label);
      in_order2 right
;;
(* val in_order2 : tree -> unit = <fun> *)

(* The first test: print the labels as they are yielded *)
let t1 = exact_reify (fun () ->
  try in_order2 tree1 
  with YT x -> Printf.printf "Got %d\n" x
);;

(* The second test: accumulate the results in a list *)
let t12 = exact_reify (fun () ->
  try in_order2 tree1; None
  with YT x -> Some x
);;

(*
val t12 : int option Ptypes.pV =
  [(0.0078125, V (Some 7)); (0.03125, V (Some 6)); (0.125, V (Some 5));
   (0.5, V (Some 4)); (0.015625, V (Some 3)); (0.25, V (Some 2));
   (0.0625, V (Some 1)); (0.0078125, V None)]
*)

(* More challenging and interesting is post_order example,
  an accumulating traversal yielding intermediate results.

  We traverse a tree post-order, yielding the running sum of
  the current label and the labels in the left and the right
  branches (that is, the sum of the labels for the subtree rooted 
  at the current node).

 Now the generator has to return a useful value.
*)

let rec post_order = function
  | Leaf -> 0
  | Node (label, left, right) ->
      let sum_left  = post_order left in
      let sum_right = post_order right in
      let sum = sum_left + sum_right + label in
      yield (YT sum); sum
;;
(* val post_order : tree -> int = <fun> *)

let tp1 = exact_reify (fun () ->
  try Printf.printf "Final: %d\n" (post_order tree1)
  with YT x -> Printf.printf "Got %d\n" x
);;

(*
Got 4
Got 5
Got 11
Got 6
Got 7
Got 16
Got 28
Final: 28
val tp1 : unit Ptypes.pV = [(1., Ptypes.V ())]
*)

(* We do not care about probabilities in this application;
   we do care about the order of results.
   The inference procedure exact_reify we used so far
   does the opposite.

  We define a new `inference' procedure, to convert the
  probabilistic model, reified into a lazy tree, into
  a stream (to flatten the tree).
*)

type 'a stream = Nil of 'a | Cons of 'a * (unit -> 'a stream);;

let rec stream_append s1 s2 =
  match s1 with
  | Nil x        -> Cons (x,s2)
  | Cons (x,s1') -> Cons (x,fun () -> stream_append (s1' ()) s2)
;;
(* val stream_append : 'a stream -> (unit -> 'a stream) -> 'a stream = <fun>
*)

(* [stream_to_list n s]
converts at most n first element of the stream to list.
[n : int option] is the limit; None means no limit.
*)

let stream_to_list n s =
  let predM = function
    | Some n -> Some (pred n)
    | None   -> None in
  let is_done = function
    | Some n when n <= 0 -> true
    | _                  -> false in
  let rec loop n acc s = 
    if is_done n then List.rev acc else
    match s () with
    | Nil x       -> List.rev (x::acc)
    | Cons (x,s)  -> loop (predM n) (x::acc) s
  in loop n [] (fun () -> s)
;;
(* val stream_to_list : int option -> 'a stream -> 'a list = <fun> *)

    
(* This function could be written more nicely.
   We will derive a more efficient implementation of generators
   anyway, see the file generator.ml
   The pattern-matching is deliberately inexhaustive;
   It covers only the cases that may occur if we use yield
   but not other Hansei operations.
*)
let stream_reify model =
  let rec loop = function
    | [(_,V x)] -> Nil x
    | [(_,C t1);(_,C t2)] -> 
	stream_append (loop (t1 ())) (fun () -> loop (t2 ()))
  in loop (reify0 model)
;;
(* val stream_reify : (unit -> 'a) -> 'a stream = <fun> *)


(* We repeat the earlier tests, using the new stream reification function *)
(* The second test: accumulate the results in a list *)
let [Some 4; Some 2; Some 5; Some 1; Some 6; Some 3; Some 7; None] = 
  stream_to_list None (stream_reify (fun () ->
  try in_order2 tree1; None
  with YT x -> Some x
));;


let _ = 
  stream_to_list None (stream_reify (fun () ->
  try Printf.printf "Final: %d\n" (post_order tree1)
  with YT x -> Printf.printf "Got %d\n" x
));;

(*
Got 4
Got 5
Got 11
Got 6
Got 7
Got 16
Got 28
Final: 28
*)