// An F# implementation of Java-style Streams
// F# has something equivalent called sequences, but I think Java streams
// are better organized.
open System
open System.Collections.Generic;;

//////// some aliases and utilities ...
type Vec<'T> = ResizeArray<'T>;;  // type alias
let omap = Option.map
let is_some = Option.isSome
let is_none = Option.isNone
let do_if = Option.iter
let unwrap_or = Option.defaultValue
let if_some f = function
  | Some(x) -> f x
  | _ -> ()
let if_else f g = function
  | Some(x) -> f(x)
  | None -> g()

///// main type implemented as a record:
type stream<'T> =
  {
     mutable generator : unit -> 'T option;  
  }
  //////////////// operations on streams
  member this.map<'U> (f: 'T -> 'U) =
    {
      stream.generator = fun () -> this.generator() |> omap f
    }
  member this.append(other:stream<'T>) = 
    let gen = this.generator;
    this.generator <- fun () ->
       let ax = gen()
       if Option.isSome ax then ax
       else other.generator()
    this
  member this.bind<'U>(bf:'T -> stream<'U>) =
    let mutable ustream = this.generator() |> omap bf
    {
       stream.generator = fun () ->
         let mutable next = None
         let mutable breakout = is_none ustream
         while not(breakout) do
           next <- ustream |> omap (fun s -> s.generator())
           if (is_none next) then
             ustream <- this.generator() |> omap bf
             breakout <- is_none ustream
           else
             breakout <- true
         //while
         next
    }
  member this.filter(p : 'T -> bool) =
    let gen = this.generator;
    this.generator <- fun () ->
      let mutable ax = gen()
      let mutable breakout = false
      while not(breakout) do
         match ax with
           | None -> breakout <- true   // stream has ended
           | Some(x) when (p x) -> breakout <- true  // keep
           | _ -> ax <- gen()                        // try next
      ax
    // this.generator
    this
  member this.limit n =
    let gen = this.generator
    let mutable counter = n
    this.generator <- fun () -> 
       if counter < 1 then None
       else (counter<-counter-1; gen())
    this
  member this.take_while (p: 'T -> bool) =
    let gen = this.generator
    this.generator <- fun strm ->
       let ax = gen()
       match ax with
         | Some x when (p x) -> ax
         | _ -> None
    this
  member this.until (p: 'T -> bool) =   // keep the last one
    let gen = this.generator
    let mutable ended = false
    this.generator <- fun strm ->
       let ax = gen()
       match (ended, ax) with
         | (false, Some x) when (p x) ->
              ended <- true
              ax
         | (false, _) -> ax
         | _ -> None
    this
  ////////////////////////////////////// running the stream
  member this.foreach (act : 'T -> unit) : unit =
    let mutable ax = this.generator()
    while Option.isSome ax do
      if_some act ax
      ax <- this.generator()
  member this.for_all (p: 'T -> bool) =
    let mutable answer = true
    let mutable ax = this.generator()
    while (ax
           |> omap (fun x -> (if not(p(x)) then answer <- false); answer)
           |> unwrap_or false) do
      ax <- this.generator();
      
    (*	    version easier to understand but uses Option.get
    while answer && Option.isSome ax do
      if not(p (Option.get ax)) then answer <- false
      else ax <- this.generator();
    *)
    
    answer
  member this.there_exists(p: 'T -> bool) = not(this.for_all(fun x -> not(p x)))
  member this.find_first(p: 'T -> bool) = // returns 'T option
    let mutable ax = this.generator();
    while (ax
           |> omap (fun x -> not(p x))
           |> unwrap_or false) do
      ax <- this.generator();
    ax      
(*
    let mutable answer = None
    while (Option.isNone answer) && (Option.isSome ax) do
      if (p (Option.get ax)) then answer <- ax
      else ax <- this.generator()
    answer
*)      
  member this.nth n =
    let mutable i = 0
    let mutable ax = this.generator()
    while (is_some ax) && i<n do
      ax <- this.generator()
      i <- i+1
    ax
  member this.get_nth n = this.nth(n) |> Option.get
  member this.find_last(p: 'T -> bool) =
    let mutable answer = None
    let mutable next = this.generator()
    while (is_some next) do
      //next |> Option.iter (fun x -> if p(x) then answer <- next)
      if unwrap_or false (omap p next) then answer <- next;
      next <- this.generator()
    answer
  member this.last() = this.find_last(fun _ -> true)
  member this.get_last() = this.last() |> Option.get


//// independent flatten function, s.bind f = flatten(s.map f)
let flatten<'T> (srcstream:stream<stream<'T>>) =
  let mutable current = srcstream.generator()
  {
     stream.generator = fun () ->
       let mutable next = None
       let mutable breakout = is_none current
       while not(breakout) do
         next <- current |> omap (fun s -> s.generator())
         if (is_none next) then
           current <- srcstream.generator()
           breakout <- is_none current
         else
           breakout <- true
       //while
       next
    }


//////// constructors

let vector_stream<'T>(v:Vec<'T>) =
  let n = v.Count
  let mutable internal_cx = 0;
  {
     generator = fun () ->
        if internal_cx < n then 
          internal_cx <- internal_cx + 1
          Some(v.[internal_cx - 1]);
        else
          None
  }

let finite_stream<'T>(v:'T list) =
  let n = v.Length
  let mutable internal_cx = 0;
  {
     generator = fun () ->
        if internal_cx < n then 
          internal_cx <- internal_cx + 1
          Some(v.[internal_cx - 1]);
        else
          None
  };;

let seeded_coinduction<'T>(seeds:'T list, gen_next: Vec<'T> -> 'T) =
  let vec = Vec<'T>()
  for x in seeds do
    vec.Add(x)
  let mutable index = 0
  {  
     generator = fun () ->
        if index < vec.Count then
           index <- index+1
           Some(vec.[index-1])
        else
           let next = gen_next vec
           vec.Add(next)
           index <- index + 1
           Some(next)
  };;

let coinduction<'T>(base_case:'T, inductive_case:'T -> 'T) =
  let mutable next = base_case
  let mutable is_first = true
  {
    stream.generator = fun () ->
      if is_first then is_first <- false
      else next <- inductive_case next
      Some(next)
  };;

// create stream by providing generator
let generate_stream<'T> (genfun:unit -> 'T) =
  {
     stream.generator = fun () -> Some(genfun());
  }

let empty_stream<'T> () = 
  {
     stream.generator = fun () -> None
  }

let option_stream opt =
  (Option.map (fun x -> finite_stream([x])) opt)
  |> unwrap_or(empty_stream())

let results_stream<'T,'E> (results:Result<'T,'E> list) =
  let rv = Vec<'T>();
  for r in results do 
    match r with
      | Ok x -> rv.Add(x)
      | _ -> ()
  vector_stream(rv)


//////////////////////////////////////// DEMO

let evens = coinduction(0, fun n -> n+2)
evens
  .filter(fun x -> x%4=0)
  .map(fun x -> float(x)/2.0)
  .until(fun x -> x>99.0)
  .foreach(printf "%.2f ")
printfn ""

let next_prime (known_primes:Vec<int>) =
  if known_primes.Count < 2 then known_primes.Count + 2 // 2 or 3
  else
    let last_prime = known_primes.[known_primes.Count - 1]
    let candidates = coinduction(last_prime+2, fun c -> c+2)
    candidates
      .find_first(fun candidate ->
          (vector_stream known_primes)
            .until(fun p -> float(p) >= 1.0+Math.Sqrt(float(candidate)))
            .for_all(fun p -> candidate % p <> 0))
      |> Option.get  // know it will exist (until overflow)

let primes known_primes = seeded_coinduction(known_primes, next_prime);

primes([]).limit(100).foreach(printf "%d ")
printfn ""

// prime factorization: collect all prime factors of a number into a vector
let prime_factors n =
  let factors = Vec<int>();
  primes([2;3;5;7;11])
    .until(fun p -> p >= n)
    .filter(fun p -> n%p=0)
    .foreach(fun f -> factors.Add(f))
  factors

let on_behalf f =
  try
    Some(f())
  with | e -> None

printf "\nEnter a number for prime factorization: "
let nopt = on_behalf (fun () -> int(Console.ReadLine()))
nopt |> if_some (fun n -> printfn "prime factors of %d: %A" n (prime_factors n))

// n!: inductive hypothesis (n,nf) means "nf is n!"
let factorial n =
  coinduction( (0,1), fun (n,nf) -> (n+1, (n+1)*nf) )
    .map(fun (_,y) -> y)
    .get_nth(n)

printfn "6! is  %d" (factorial 6)

// 2x2 matrix multiplication
type M2x2 = (int64*int64*int64*int64)

let mmult(A:M2x2) (B:M2x2) = 
  let ((a, b,
        c, d), (q, r,
                s, t)) = (A,B)
  (a*q+b*s, a*r+b*t,
   c*q+d*s, c*r+d*t)

// using int64 types - L suffix required
let IDM = (1L, 0L,
           0L, 1L)
let FIBM = (1L, 1L,
            1L, 0L)

let power(basis, exponent, multiplier, identity) =  // "base" is a keyword
   let next(ax,factor,n) =
     let newfactor = multiplier factor factor
     ((if n%2=1 then (multiplier ax factor) else ax), newfactor, n/2)
   coinduction((identity, basis, exponent), next)
      .until(fun (_,_,e) -> e=0)
      .map(fun(ax,_,_) -> ax)
      .get_last()

printfn "2**9 is %d" ( power(2,9,(*),1) )

// nth Fibonacci number in log n steps:
let Fib n  =
  let (a,_,_,_) = power(FIBM, n-1, mmult, IDM)
  a

printfn "Fib(9): %A"  (Fib 9)
printfn "Fib(10): %A" (Fib 10)
printfn "Fib(11): %A" (Fib 11)
printfn "Fib(100): %A" (Fib 100)