open Inf

let unimplemented () = failwith "unimplemented";;

(* an infinite stream of zeros *)
let zeros = const 0;;

(* the stream of all natural numbers (represented as integers, cycling) *)
let naturals : int stream = 
  scan (fun base x -> base+1) (-1) zeros 
;;

(* all even natural numbers *)
let evens : 'a stream = 
  map (fun v -> v*2) naturals
;;

(* all odd natural numbers *)
let odds : 'a stream = 
  map (fun v -> v+1) evens
;;

let filter_multiples (n:int) (s:int stream) : int stream =
  filter (fun x -> not (x mod n == 0)) s
;;

(* all prime numbers *)
let primes : 'a stream = 
  let rec prime_loop s = 
    let (v,s') = next s in
    cons v (fun () -> prime_loop (filter_multiples v s'))
  in
  let (_,s1) = next naturals in
  let (_,s2) = next s1 in
  prime_loop s2
;;

(***********)
(* Testing *)
(***********)

let rec take (n:int) (s:'a stream) : 'a list =
  let rec aux n s result =
    if n <= 0 then 
      result
    else
      let (x,s') = next s in
      aux (n-1) s' (x::result)
  in
  List.rev (aux n s [])
;;

let l1 = take 10 naturals;;
let l2 = take 11 evens;;
let l3 = take 8 odds;;
let l4 = take 7 primes;;
let ls = [l1;l2;l3;l4];;

let print_list l =
  let rec print_elements l =
    match l with
	[] -> ()
      | [x] -> print_int x
      | x::tail -> print_int x; print_string ", "; print_elements tail
  in
  print_string "[";
  print_elements l;
  print_string "]\n"
;;

let main () =
  List.map print_list ls;
;;

main ();;