open Inf

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

(* the stream of all natural numbers (represented as integers, cycling) *)
let naturals : int stream = 
  scan (+) (-1) (const 1)

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



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



(* F_n = F_n-1 + F_n-2 *)
let fibs =
  let pairs = scan (fun (f2,f1) x -> (f1, f1+f2)) (0,1) zeros in
  cons 0 (fun () -> map (fun (n1,n2) -> n1) pairs)


(* return a stream like s but containing only those numbers x such
   that x mod n != 0 *)
let filter_multiples (n:int) (s:int stream) : int stream =
  filter (fun x -> (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 todo () = failwith "todo"

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 
    [take 10 fibs;  (* change me! *)
     take 50 primes]
;;

main ();;