streamIO.scala

trait Functor[F[_]] {
  def map[A,B](fa: F[A])(f: A => B): F[B]

  def distribute[A,B](fab: F[(A, B)]): (F[A], F[B]) =
    (map(fab)(_._1), map(fab)(_._2))

  def codistribute[A,B](e: Either[F[A], F[B]]): F[Either[A, B]] = e match {
    case Left(fa) => map(fa)(Left(_))
    case Right(fb) => map(fb)(Right(_))
  }
}

object Functor {
  val listFunctor = new Functor[List] {
    def map[A,B](as: List[A])(f: A => B): List[B] = as map f
  }
}

trait Monad[F[_]] extends Functor[F] {
  def unit[A](a: => A): F[A]

  def flatMap[A,B](ma: F[A])(f: A => F[B]): F[B]

  def map[A,B](ma: F[A])(f: A => B): F[B] =
    flatMap(ma)(a => unit(f(a)))
  def map2[A,B,C](ma: F[A], mb: F[B])(f: (A, B) => C): F[C] =
    flatMap(ma)(a => map(mb)(b => f(a, b)))

  def sequence[A](lma: List[F[A]]): F[List[A]] =
    lma.foldRight(unit(List[A]()))((ma, mla) => map2(ma, mla)(_ :: _))

  def traverse[A,B](la: List[A])(f: A => F[B]): F[List[B]] =
    la.foldRight(unit(List[B]()))((a, mlb) => map2(f(a), mlb)(_ :: _))

  def join[A](mma: F[F[A]]): F[A] = flatMap(mma)(ma => ma)
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

/**
 * A Process[I,O] can be used to transform a stream containing I values to a stream of O values.
 */
sealed trait Process[I,O] {
  def apply(s: Stream[I]): Stream[O] = this match {
    case Halt() => Stream()
    case Await(recv) => s match {
      case h #:: t => recv(Some(h))(t)
      case xs => recv(None)(xs)
    }
    case Emit(h,t) => h #:: t(s)
  }

  def repeat: Process[I,O] = {
    def go(p: Process[I,O]): Process[I,O] = p match {
      case Halt() => go(this)
      case Await(recv) => Await {
        case None => recv(None)
        case i => go(recv(i))
      }
      case Emit(h, t) => Emit(h, go(t))
    }
    go(this)
  }

  def |>[O2](p2: Process[O,O2]): Process[I,O2] = p2 match {
    case Halt() => Halt()
    case Emit(h, t) => Emit(h, this |> t)
    case Await(recv) => this match {
      case Halt() => Halt()
      case Await(recv2) => await( (i: I) => recv2(Some(i)) |> p2 )
      case Emit(h, t) => t |> recv(Some(h))
    }
  }

  def ++(p: => Process[I,O]): Process[I,O] = this match {
    case Halt() => p
    case Emit(h, t) => Emit(h, t ++ p)
    case Await(recv) => Await(recv andThen (_ ++ p))
  }

  def flatMap[O2](f: O => Process[I,O2]): Process[I,O2] = this match {
    case Halt() => Halt()
    case Emit(h, t) => f(h) ++ t.flatMap(f)
    case Await(recv) => Await(recv andThen (_ flatMap f))
  }

}

// indicates to the driver that the head value should be emitted to the output stream, and the machine should then be
// transitioned to the tail state.
case class Emit[I,O](head: O, tail: Process[I,O] = Halt[I,O]()) extends Process[I,O]

// requests a value from the input stream. The driver should pass the next available value to the recv function, or
// None if the input has no more elements.
case class Await[I,O](recv: Option[I] => Process[I,O]) extends Process[I,O]

// indicates to the driver that no more elements should be read from the input or emitted to the output.
case class Halt[I,O]() extends Process[I,O]


def liftOne[I,O](f: I => O): Process[I,O] = Await {
    case Some(i) => Emit(f(i))
    case None => Halt()
  }

def lift[I,O](f: I => O): Process[I,O] = liftOne(f).repeat

def emit[I,O](head: O, tail: Process[I,O] = Halt[I,O]()): Process[I,O] =  Emit(head, tail)

lift((i: Int) => s"value $i")(Stream(1, 2, 3))

def count[I]: Process[I,Int] = {
  def go(i: Int): Process[I, Int] = Await[I, Int] {
    case Some(_)   => Emit(i + 1, go(i + 1))
    case None => Halt()
  }
  go(0)
}

def await[I, O](recv: I => Process[I,O]): Await[I, O] = Await {
  case Some(i) => recv(i)
  case None => Halt()
}

def loop[S,I,O](z: S)(f: (I,S) => (O,S)): Process[I,O] =
  await((i: I) => f(i,z) match {
    case (o,s2) => emit(o, loop(s2)(f))
  })

def monad[I]: Monad[({ type f[x] = Process[I,x]})#f] =
  new Monad[({ type f[x] = Process[I,x]})#f] {
    def unit[O](o: => O): Process[I,O] = Emit(o)
    def flatMap[O,O2](p: Process[I,O])(f: O => Process[I,O2]): Process[I,O2] = p flatMap f
  }