import scala.language.higherKinds
import scala.language.implicitConversions

object monadHelper {

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

  // TODO: use standard monad library

  trait Functor[M[_]] {
    def map[A, B](self: M[A], f: A => B): M[B]
  }

  trait Monad[M[_]] extends Functor[M] {
    def unit[A](value: A): M[A]
    def flatMap[A, B](self: M[A], f: A => M[B]): M[B]
  }

  implicit class MonadInterface[M[_], A](val self: M[A]) extends AnyVal {
    def map[B](f: A => B)(implicit m: Functor[M]): M[B] = m.map(self, f)
    def flatMap[B](f: A => M[B])(implicit m: Monad[M]): M[B] = m.flatMap(self, f)
    def foreach(f: A => Unit)(implicit m: Functor[M]) = map(f)
  }

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

}

object lazyMonad {

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

  // minimal operational monad

  sealed trait Opr[F[_], A]
  case class Pure[F[_], A](value: A) extends Opr[F, A]
  case class Bind[F[_], A, B](head: F[A], tail: A => Opr[F, B]) extends Opr[F, B]

  object Opr {
    def apply[F[_], A](value: A): Opr[F, A] = Pure(value)
    def lift[F[_], A](from: F[A]): Opr[F, A] = bind(from) { (res: A) => Pure(res) }
    def bind[F[_], A, B](head: F[A])(tail: A => Opr[F, B]): Opr[F, B] = Bind(head, tail)
  }

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

  // Lazy implementation

  sealed trait Lazy[+A] {
    def forceMemo: Option[A]

    def isForced: Boolean = forceMemo.isDefined

    def force: A = forceMemo match {
      case Some(memo) => memo
      case None => {
        var cont: Opr[Lazy, A] = Opr.lift(this)
        while (true) {
          cont match {
            case Pure(value) => return value
            case Bind(head, tail) => head.forceMemo match {
              case None       => cont = head.forceImpl(tail)
              case Some(memo) => cont = tail(memo)
            }
          }
        }
        ???  // TODO: report proper exception
      }
    }

    def forceImpl[R](cont: A => Opr[Lazy, R]): Opr[Lazy, R]

    override def toString: String = {
      if (isForced) {
        s"Lazy(${force.toString})"
      } else {
        "Lazy(#)"
      }
    }
  }

  case class SimpleLazy[A](promise: () => A) extends Lazy[A] {
    var forceMemo: Option[A] = None

    def forceImpl[R](cont: A => Opr[Lazy, R]): Opr[Lazy, R] = {
      val result = promise()
      forceMemo = Some(result)
      cont(result)
    }
  }

  case class MappedLazy[A, B](head: Lazy[A], tail: A => Lazy[B]) extends Lazy[B] {
    var forceMemo: Option[B] = None

    def forceImpl[R](cont: B => Opr[Lazy, R]): Opr[Lazy, R] = {
      Opr.bind(head) { (mid: A) =>
        Opr.bind(tail(mid)) { (res: B) =>
          forceMemo = Some(res)
          cont(res)
        }
      }
    }
  }

  object Lazy {
    def apply[A](promise: => Lazy[A]): Lazy[A] = SimpleLazy(() => promise).join
    def unit[A](value: => A): Lazy[A] = SimpleLazy(() => value)
  }

  implicit class LazyInterface[A](val self: Lazy[A]) extends AnyVal {
    def map[B](f: A => B): Lazy[B] = flatMap { (a: A) =>
      val b: B = f(a)
      Lazy.unit(b)
    }

    def flatMap[B](f: A => Lazy[B]): Lazy[B] = MappedLazy(self, f)
  }

  // implicit def makeLazy[A](promise: => A): Lazy[A] = SimpleLazy(() => promise)

  trait LazyFamily[A] {
    def join(nested: Lazy[A]): A
  }

  implicit def joinImplForLazy[A]: LazyFamily[Lazy[A]] = new LazyFamily[Lazy[A]] {
    def join(nested: Lazy[Lazy[A]]): Lazy[A] = nested.flatMap(a => a)
  }

  implicit class LazyInterface2[A](val self: Lazy[A]) extends AnyVal {
    def join(implicit impl: LazyFamily[A]): A = impl.join(self)
  }

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

}

object lazyList {
  // usage example

  import lazyMonad._
  import monadHelper._

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

  implicit val lazyMonadInstance: Monad[Lazy] = new Monad[Lazy] {
    def unit[A](value: A) = Lazy.unit(value)
    def map[A, B](self: Lazy[A], f: A => B): Lazy[B] = self.map(f)
    def flatMap[A, B](self: Lazy[A], f: A => Lazy[B]): Lazy[B] = self.flatMap(f)
  }

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

  trait ListF[+A, +R]
  case object NullF extends ListF[Nothing, Nothing]
  case class ConsF[A, R](head: A, tail: R) extends ListF[A, R]

  class LazyList[+A](val impl: Lazy[ListF[A, LazyList[A]]]) /* extends AnyVal */ {
    override def toString: String = impl.toString
  }

  val LazyNull: LazyList[Nothing] = new LazyList(Lazy.unit(NullF))

  implicit def joinImplForLazyList[A]: LazyFamily[LazyList[A]] = new LazyFamily[LazyList[A]] {
    def join(nested: Lazy[LazyList[A]]): LazyList[A] = new LazyList(nested.flatMap(_.impl))
  }

  implicit class LazyListInterface[A](val self: LazyList[A]) extends AnyVal {
    def demolish[R](caseNull: => R)(caseCons: (A, LazyList[A]) => R): Lazy[R] = {
      self.impl.map { (impl: ListF[A, LazyList[A]]) =>
        impl match {
          case NullF             => caseNull
          case ConsF(head, tail) => caseCons(head, tail)
        }
      }
    }
    def #::[AA >: A](head: => AA): LazyList[AA] = new LazyList(Lazy.unit(ConsF(head, self)))
    def isEmpty: Boolean = demolish(true) { (_: A, _: LazyList[A]) => false }.force
    def isNonEmpty: Boolean = !isEmpty
    def head: A = demolish(???) { (h: A, _: LazyList[A]) => h }.force
    def tail: LazyList[A] = demolish(???) { (_: A, t: LazyList[A]) => t }.force
    def foldRight[R](init: R)(f: (A, R) => R): Lazy[R] = self.demolish(Lazy.unit(init)) { (h: A, t: LazyList[A]) => t.foldRight(init)(f).map(tt => f(h, tt)) }.flatMap(a => a)
    def zip[B](that: LazyList[B]): LazyList[(A, B)] = demolish[LazyList[(A, B)]](LazyNull) { (h1: A, t1: LazyList[A]) => that.demolish[LazyList[(A, B)]](LazyNull) { (h2: B, t2: LazyList[B]) => (h1, h2) #:: t1.zip(t2) }.join }.join
    def map[B](f: A => B): LazyList[B] = demolish[LazyList[B]](LazyNull) { (h: A, t: LazyList[A]) => f(h) #:: new LazyListInterface(t).map(f) }.join
    def take(n: Int): LazyList[A] = {
      if (n <= 0) { LazyNull } else { demolish[LazyList[A]](LazyNull) { (h: A, t: LazyList[A]) => h #:: t.take(n - 1) }.join }
    }
  }

  object LazyList {
    def apply[A](elems: Lazy[A]*): LazyList[A] = {
      lazy val loop: (Int) => LazyList[A] = { (i: Int) =>
        if (i >= elems.length) {
          LazyNull
        } else {
          elems(i).map(_ #:: loop(i + 1)).join
        }
      }
      loop(0)
    }
  }

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

}

object main {
  import lazyMonad._
  import lazyList._

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

  lazy val fibs: LazyList[Int] = 0 #:: 1 #:: Lazy.unit {
    fibs.zip(fibs.tail).map { case (a, b) => a + b }
  }.join

  lazy val nats: LazyList[Int] = 0 #:: Lazy.unit(nats.map(_ + 1)).join

  def main(args: Array[String]) = {
    // calculations are cached
    println(s"fibs: $fibs")
    println(s"fibs.take(20).sum: ${fibs.take(20).foldRight(0)(_ + _).force}")
    println(s"fibs: $fibs")
    // deep recursion does not cause stack overflow
    println(s"nats.take(10^6).sum: ${nats.take(1000000).foldRight(0)(_ + _).force}")
  }

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

}