(** Extensions to the Stream module. *)

(** Return a stream that yields only the elements that satisfy PRED. *)
let filter pred s =
  Stream.from
    (fun n ->
       let rec loop () =
	 match Stream.peek s with
	   Some x ->
	     Stream.junk s;
	     if pred x
	     then Some x
	     else loop ()
	 | None -> None in
       loop ())

(** Return a list of the stream elements *)
let to_list s =    
  let result = ref [] in
  Stream.iter (fun x -> result := x :: !result) s;
  List.rev !result

let map f s =
  Stream.from (fun n ->
		 match Stream.peek s with
		   Some x ->
		     Stream.junk s;
		     Some (f x)
		 | None -> None)

(*let map_test () =
  Stream.iter
    (fun n -> print_endline (string_of_int n))
    (map (fun x -> x + x) (Stream.of_list [1;2;3;4]))*)

(* Equivalent to:

  Stream.map
    Option.get
    (Stream.filter
       Option.pred
       (Stream.map f s)) *)

let map_filter (f : ('a -> 'b option)) (s : 'a Stream.t) : ('b Stream.t) =
  Stream.from (fun n ->
		 let rec loop () =
		   match Stream.peek s with
		     Some x -> begin
		       Stream.junk s;
		       match f x with
			 Some y -> Some y
		       | None -> loop ()
		     end
		   | None -> None in
		 loop ())

(*let map_filter_test () =
  Stream.iter
    (fun n -> print_endline (string_of_int n))
    (map_filter
       (fun x -> if x < 3 then Some x else None)
       (Stream.of_list [1;2;3;4]))*)

(* Like map, but pass the option to f.  This way we get a call to f
   following the last element -- useful for closing channels. *)
let mapopt f s =
  Stream.from (fun n ->
		 try f (Some (Stream.next s))
		 with Stream.Failure -> f None)

(** Junk n elements *)
let njunk n s = for i = 1 to n do Stream.junk s done

(** Turn a stream of streams of x into a stream of x. *)
let flatten (s : 'a Stream.t Stream.t) : 'a Stream.t =
  Stream.from
    (fun n ->
       let rec loop () =
	 match Stream.peek s with
	   Some t -> begin
	     match Stream.peek t with
	       Some e -> 
		 Some (Stream.next t)
	     | None ->
		 Stream.junk s;
		 loop ()
	   end
	 | None -> None in
       loop ())

(*let flatten_test () =
  let s = Stream.of_list [Stream.of_list [1;2;3;4];
			  Stream.of_list [5;6;7];
			  Stream.of_list [8; 9]] in
  let t = flatten s in
  Stream.iter (fun n -> print_endline (string_of_int n)) t*)

(** Return a stream that concatenates a list of streams (deprecated) *)
let concat lst = flatten (Stream.of_list lst)

(* Split a stream into substreams when the predicate matches *)
let chop (pred : 'a -> bool) (s : 'a Stream.t) : ('a Stream.t Stream.t) =
  Stream.from
    (fun _ ->
       match Stream.peek s with
	 None -> None
       | Some x ->
	   Some (Stream.from
		   (fun n ->
		      match Stream.peek s with
			Some x when (n = 0 || not (pred x)) ->
			  (* The first element that matches the predicate
			     is part of the stream, or any element that does
			     not match the predicate. *)
			  Stream.junk s;
			  Some x
		      | _ ->
			  None)))

(*let chop_test () =
  let s =
    chop (fun x -> x / 2 * 2 = x) (Stream.of_list [1;3;4;5;7;2;6;5;9]) in
  Stream.iter 
    (fun t ->
       Stream.iter (fun n -> print_string (" " ^ string_of_int n)) t;
       print_endline "") s*)

(** Deprecated: Turn a stream of 'a option into a stream of streams of 'a.
    use chop. *)
let split (s : 'a option Stream.t) : ('a Stream.t Stream.t) =
  let t = chop (function None -> true | _ -> false) s in
  map (map_filter (fun x -> x)) t

(*let split_test s =
  let t = split (Stream.of_list [Some 1; Some 2; None; Some 3;
				 None; None; Some 5]) in
  Stream.iter
    (fun u ->
       print_endline ("["
		      ^ String.concat "; " (List.map string_of_int (to_list u))
		      ^ "]")) t*)

let sort cmp s =
  Stream.of_list (List.sort cmp (to_list s))