(* Lecture 14: Infinite data structures *)

type 'a stream = Cons of 'a * 'a stream

let head (Cons(x, _) : 'a stream) : 'a        = x
let tail (Cons(_, r) : 'a stream) : 'a stream = r

let rec ones = Cons(1, ones)

assert( head ones = 1 )
assert( head (tail ones) = 1 )
assert( head (tail (tail ones)) = 1 )

assert( phys_equal ones (tail ones) )
assert( phys_equal ones (tail (tail ones)) )

let rec count_from (n : int) : int stream =
  Cons(n, count_from (n + 1))

(* This is gonna blow the stack: *)
let rec nats = count_from 0

let rec map (Cons(x, r) : 'a stream) ~(f: 'a -> 'b) : 'b stream =
  Cons(f x, map r ~f)

(* Not allowed by OCaml: *)
(*
let rec nats = Cons(1, map ~f:((+) 1) nats)
*)

(* Diverges: *)
let twos = map ~f:((+) 1) ones

(** SECOND ATTEMPT: delay computing the tail **)

(* Doesn't type check: *)
(*
let rec map (Cons(x, r)) ~f =
  Cons(f x, fun () -> map r ~f)
 *)

type 'a stream = unit -> 'a str
 and 'a str = Cons of 'a * 'a stream

type 'a stream = unit -> 'a str
 and 'a str = { hd : 'a; tl : 'a stream }

let rec ones : int stream = fun () -> { hd = 1; tl = ones }

let head (s : 'a stream) : 'a = (s ()).hd

assert( head ones = 1 )

let tail (s : 'a stream) : 'a stream = (s ()).tl

assert( head (tail ones) = 1 )
assert( head (tail (tail ones)) = 1 )
assert( phys_equal ones (tail ones) )
assert( phys_equal (tail ones) (tail ones) )

let rec map (s : 'a stream) ~(f : 'a -> 'b) : 'b stream =
  fun () -> { hd = f (head s); tl = map (tail s) ~f }

let twos = map ones ~f:((+) 1)

assert( head twos = 2 )
assert( head (tail twos) = 2 )
assert( head (tail (tail twos)) = 2 )
assert( not (phys_equal twos (tail twos)) )
assert( not (phys_equal (tail twos) (tail twos)) )

let rec take (s : 'a stream) (n : int) : 'a list =
  if n > 0 then head s :: take (tail s) (n - 1)
  else []

assert( take twos 5 = [2; 2; 2; 2; 2] )

let rec count_from (n : int) : int stream =
  fun () -> { hd = n; tl = count_from (n + 1) }

let nats = count_from 0

assert( take nats 5 = [0; 1; 2; 3; 4] )

(* Or: *)
let rec nats = fun () -> { hd = 0; tl = map ~f:((+) 1) nats }

assert( take nats 5 = [0; 1; 2; 3; 4] )


(* Map over two streams item-wise: *)
let rec map2 (s1 : 'a stream) (s2 : 'b stream)
             ~(f : 'a -> 'b -> 'c) : 'c stream =
  fun () -> { hd = f (head s1) (head s2); tl = map2 (tail s1) (tail s2) ~f }

assert( take (map2 ~f:(+) ones twos) 5 = [3; 3; 3; 3; 3] )

let rec fibs =
  fun () -> { hd = 0;
              tl = fun () -> { hd = 1;
                               tl = map2 ~f:(+) fibs (tail fibs) } }

assert( take fibs 10 = [0; 1; 1; 2; 3; 5; 8; 13; 21; 34] )

let _ = take fibs 17

(* Exponential time? *)

(** THIRD ATTEMPT: memoize **)

type 'a promise = Unevaluated of (unit -> 'a) | Evaluated of 'a

type 'a stream = 'a str promise ref
 and 'a str = { hd : 'a; tl : 'a stream }

let rec ones = ref (Unevaluated (fun () -> { hd = 1; tl = ones }))

let rec head (s : 'a stream) : 'a =
  match !s with
  | Evaluated str -> str.hd
  | Unevaluated f ->
     let str = f () in
     s := Evaluated str;
     str.hd

let rec tail (s : 'a stream) : 'a stream =
  match !s with
  | Evaluated str -> str.tl
  | Unevaluated f ->
     let str = f () in
     s := Evaluated str;
     str.tl

type 'a my_lazy_t = 'a promise ref

let force (pr : 'a my_lazy_t) : 'a = match !pr with
  | Evaluated v -> v
  | Unevaluated f ->
     let v = f () in
     pr := Evaluated v;
     v

type 'a stream = 'a str my_lazy_t
 and 'a str = { hd : 'a; tl : 'a stream }

let rec ones = ref (Unevaluated (fun () -> { hd = 1; tl = ones }))

let rec head (s : 'a stream) : 'a = (force s).hd

assert( head ones = 1 )

let rec tail (s : 'a stream) : 'a stream = (force s).tl

assert( head (tail ones) = 1 )
assert( head (tail (tail ones)) = 1 )

let rec take (s : 'a stream) (n : int) : 'a list =
  if n > 0
  then head s :: take (tail s) (n - 1)
  else []

assert( take ones 5 = [1; 1; 1; 1; 1] )

let rec map (s : 'a stream) ~(f : 'a -> 'b) : 'b stream =
  ref (Unevaluated (fun () ->
                    { hd = f (head s); tl = map (tail s) ~f }))

let twos = map ~f:((+) 1) ones

assert( take twos 5 = [2; 2; 2; 2; 2] )

(** FOURTH (FINAL) ATTEMPT **)

type 'a stream = 'a str lazy_t
 and 'a str = { hd : 'a; tl : 'a stream }

let head (s : 'a stream) : 'a = (Lazy.force s).hd

let tail (s : 'a stream) : 'a stream = (Lazy.force s).tl

let rec map (s : 'a stream) ~(f : 'a -> 'b) : 'b stream =
  lazy { hd = f (head s);
         tl = map (tail s) ~f }

let rec take (s : 'a stream) (n : int) : 'a list =
  if n > 0
  then head s :: take (tail s) (n - 1)
  else []

let rec ones = lazy { hd = 1; tl = ones }

assert( take ones 5 = [1; 1; 1; 1; 1] )

let rec count_from (n : int) : int stream =
  lazy { hd = n;
         tl = count_from (n + 1) }

let nats = count_from 0

assert( take nats 5 = [0; 1; 2; 3; 4] )

let rec nats = lazy { hd = 0; tl = map ~f:((+) 1) nats }

assert( take nats 5 = [0; 1; 2; 3; 4] )

let rec map2 (s1 : 'a stream) (s2 : 'b stream)
             ~(f : 'a -> 'b -> 'c) : 'c stream =
  lazy { hd = f (head s1) (head s2);
         tl = map2 (tail s1) (tail s2) ~f }

let rec fibs = lazy { hd = 0;
                      tl = lazy { hd = 1;
                                  tl = map2 ~f:(+) fibs (tail fibs) } }

assert( take fibs 10 = [0; 1; 1; 2; 3; 5; 8; 13; 21; 34] )

let _ = take fibs 40

(* This version is insufficiently lazy: *)
let rec filter (s : 'a stream) ~(f : 'a -> bool) : 'a stream =
  if f (head s)
  then lazy { hd = head s;
              tl = filter (tail s) ~f }
  else filter (tail s) ~f

let rec filter (s : 'a stream) ~(f : 'a -> bool) : 'a stream =
  let rec loop r =
    if f (head r)
    then { hd = head r;
           tl = lazy (loop (tail r)) }
    else loop (tail r) in
  lazy (loop s)

let even n = n mod 2 = 0
let odd n = not (even n)

let evens = filter nats ~f:even

assert( take evens 5 = [0; 2; 4; 6; 8] )

let odds = filter nats ~f:odd

assert( take odds 5 = [1; 3; 5; 7; 9] )

(* Sieve of Eratosthenes *)

let not_div_by m n = n mod m <> 0

let rec sieve (s : int stream) : int stream =
  lazy { hd = head s;
         tl = sieve (filter (tail s) ~f:(not_div_by (head s))) }

let primes = sieve (tail (tail nats))

assert( take primes 10 = [2; 3; 5; 7; 11; 13; 17; 19; 23; 29] )

(* Approximating Taylor series *)

(* e = 1/0! + 1/1! + 1/2! + 1/3! + ... *)

let rec fact_f n = if n <= 0 then 1.0 else Float.of_int n *. fact_f (n - 1)

let e_series = map2 ~f:(/.) (map ~f:Float.of_int ones)
                            (map ~f:fact_f nats)

let sum_up_to_n : float stream -> int -> float =
  let rec loop acc s n =
    if n <= 0
    then acc
    else loop (acc +. head s) (tail s) (n - 1) in
  loop 0.0

let _ = sum_up_to_n e_series 1
let _ = sum_up_to_n e_series 5
let _ = sum_up_to_n e_series 10
let _ = sum_up_to_n e_series 100
let _ = sum_up_to_n e_series 1000

(* pi = 4/1 - 4/3 + 4/5 - 4/7 + 4/9 - ... *)

let rec alt_fours = lazy { hd = 4.;
                           tl = lazy { hd = -4.;
                                       tl = alt_fours } }

let pi_series = map2 ~f:(/.) alt_fours (map ~f:Float.of_int odds)

let _ = sum_up_to_n pi_series 1
let _ = sum_up_to_n pi_series 5
let _ = sum_up_to_n pi_series 10
let _ = sum_up_to_n pi_series 100
let _ = sum_up_to_n pi_series 1000
let _ = sum_up_to_n pi_series 10000
let _ = sum_up_to_n pi_series 100000

(* Automatic numeric integration *)

let trap_area (f : float -> float) (a : float) (b : float) : float =
  (f a +. f b) *. (b -. a) /. 2.

let mid (a : float) (b : float) : float =
  (a +. b) /. 2.

let rec integrate (f : float -> float) (a : float) (b : float) : float stream =
  lazy { hd = trap_area f a b;
         tl = map2 ~f:(+.) (integrate f a (mid a b))
                           (integrate f (mid a b) b) }

let rec within eps s =
  if Float.abs (head s -. head (tail s)) < eps
  then head s
  else within eps (tail s)

let integral (f : float -> float) (a : float) (b : float) (eps : float) : float
  = within eps (integrate f a b)

let _ = integral (fun x -> x *. x) 0. 1. 0.1
let _ = integral (fun x -> x *. x) 0. 1. 0.01
let _ = integral (fun x -> x *. x) 0. 1. 0.001
let _ = integral (fun x -> x *. x) 0. 1. 0.0000001