(* CS320 Fall 2006 Assignment 5.
 * Likai Liu, Computer Science Department, Boston University.
 *
 * Distribution without permission is prohibited and is a violation of U.S.
 * Copyright Law. Submitting this solution as your own work whether in part
 * or as a whole also violates Academic Code of Conduct.
 *)

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

(* General purpose stream processors *)

fun from n = Cons (n, fn () => from (n + 1))

fun map (f: 'a -> 'b) (xs: 'a stream): 'b stream =
    case xs of
      Nil => Nil
    | Cons (x, xS) => Cons (f x, fn () => map f (xS ()))

fun mapAccu (f: 'a * 'c -> 'b * 'c) (init: 'c) (xs: 'a stream): 'b stream =
    case xs of
      Nil => Nil
    | Cons (x, xS) =>
      let val (y, res) = f (x, init) in
	Cons (y, fn () => mapAccu f res (xS ()))
      end

fun filter (p: 'a -> bool) (xs: 'a stream): 'a stream =
  case xs of
      Nil => Nil
    | Cons (x, xS) =>
      if p x then Cons (x, fn () => filter p (xS ())) else filter p (xS ())

fun hd Nil = raise Empty
  | hd (Cons (x, _)) = x

fun tl Nil = raise Empty
  | tl (Cons (_, xS)) = xS ()

fun peek (xs: 'a stream): 'a option =
    SOME (hd xs) handle Empty => NONE

fun npeek n xs =
    let
      fun loop 0 _ accu = List.rev accu
	| loop n xs accu =
	  (case peek xs of
	     NONE => loop 0 xs accu
	   | SOME x => loop (n - 1) (tl xs) (x :: accu))
    in
      if n < 0 then raise Subscript else loop n xs []
    end

fun repeat 0 f x = x
  | repeat n f x = repeat (n - 1) f (f x)

fun skip n xs = repeat n tl xs
fun nth n xs = hd (skip n xs)

(* Euler Transform from class (almost) *)

fun euler xs =
    case npeek 3 xs of
      [x0, x1, x2] =>
      let val x01 = x0 - x1 and x21 = x2 - x1 in
	Cons (x2 - x21 * x21 / (x21 + x01), fn () => euler (tl xs))
      end
    | _ => Nil

(* Problem 1 *)

fun partialSum (xs: real stream) =
    mapAccu (fn (x, accu) => let val sum = x + accu in (sum, sum) end) 0.0 xs

fun alternateNeg (xs: real stream) =
    mapAccu (fn (x, flag) => if flag then (x, false) else (~x, true)) true xs

fun recip (xs: real stream) = map (fn x => 1.0 / x) xs

fun intsToReals (xs: int stream) = map Real.fromInt xs

val ln2 = partialSum (alternateNeg (recip (intsToReals (from 1))))

val ln2_fast = repeat 8 euler ln2

(* Problem 2 *)

fun sieve (xs: int stream): int stream =
  let
     val Cons (x, xS) = xs
     val xs = filter (fn y  => y mod x > 0) (xS ())
  in
     Cons (x, fn () => sieve xs)
  end

val primes = sieve (from 2)

val p2 = partialSum (recip (intsToReals primes))

(* Problem 3 *)

fun merge (leq: 'a * 'a -> bool) (xs: 'a stream, ys: 'a stream): 'a stream =
    case (xs, ys) of
      (Nil, _) => ys
    | (_, Nil) => xs
    | (Cons (x, xS), Cons (y, yS)) =>
      if leq (x, y) then
	Cons (x, fn () => merge leq (xS (), ys))
      else
	Cons (y, fn () => merge leq (xs, yS ()))

local
  (* open IntInf *)
  fun from n = Cons (n, fn () => from (n + 1))
  fun cube x = x * x * x
  fun cubesum (i, j) = cube i + cube j
in
  (* This version only enumerates upper diagonal so we can avoid trivial,
   * non-distinctive pairs.
   *)

  fun sortedPairs (n: int): (int * int) stream =
      Cons ((n, n),
	    fn () =>
	       merge
		 (fn (p1, p2) => (cubesum p1 <= cubesum p2))
		 (map (fn x => (n, x)) (from (n + 1)), sortedPairs (n + 1)))

  fun ramanujanPairs () =
      let
	fun collectPairs (xs, n, len, accu) =
	    case peek xs of
	      NONE =>
	      if len > 1 then Cons ((n, List.rev accu), fn () => Nil) else Nil
	    | SOME p =>
	      if cubesum p = n then
		collectPairs (tl xs, n, len + 1, p :: accu)
	      else
		if len > 1 then
		  Cons ((n, List.rev accu),
			fn () => collectPairs (tl xs, cubesum p, 1, [p]))
		else
		  collectPairs (tl xs, cubesum p, 1, [p]);
      in
	case sortedPairs 0 of
	  Nil => Nil
	| Cons (p, xS) => collectPairs (xS (), cubesum p, 1, [p])
      end
end

val ramanujan = map #1 (ramanujanPairs ())