code.scala

import scala.annotation.tailrec

trait RNG {
  def nextInt: (Int, RNG)
}

case class SimpleRNG(seed: Long) extends RNG {
  def nextInt: (Int, RNG) = {
    val newSeed = (seed * 0x5DEECE66DL + 0xBL) & 0xFFFFFFFFFFFFL
    val nextRNG = SimpleRNG(newSeed)
    val n = (newSeed >>> 16).toInt
    (n, nextRNG)
  }
}

object RNG {
  type Rand[+A] = RNG => (A, RNG)

  def unit[A](a: A): Rand[A] =
    rng => (a, rng)

  def map[A,B](s: Rand[A])(f: A => B): Rand[B] =
    rng => {
      val (a, rng2) = s(rng)
      (f(a), rng2)
    }

  def map2[A,B,C](ra: Rand[A], rb: Rand[B])(f: (A, B) => C): Rand[C] =
    rng => {
      val (a, rnga2) = ra(rng)
      val (b, rngb2) = rb(rnga2)
      (f(a, b), rng)
    }

  def nonNegativeInt(): Rand[Int] =
    map(r => r.nextInt)(i => i.abs)

  def double(): Rand[Double] =
    map(nonNegativeInt())(i => i.toDouble / Int.MaxValue)

  def sequence[A](fs: List[Rand[A]]): Rand[List[A]] =
    fs.foldRight(unit(Nil) : Rand[List[A]])((a, z) => map2(a, z)(_ :: _))

  def flatMap[A,B](f: Rand[A])(g: A => Rand[B]): Rand[B] =
    rng => {
      val (a, rng2) = f(rng)
      g(a)(rng2)
    }

  @tailrec
  def nonNegativeLessThan(n: Int): Rand[Int] =
    flatMap(nonNegativeInt())(a => if (a < n) { unit(a) } else { nonNegativeLessThan(n)} )
}

object Hello extends App {
  val rng = SimpleRNG(42)
  val rngList = List.fill(10)((r: RNG) => RNG.double()(r))
  println(rng.nextInt)
  println(RNG.double()(rng))
  println(RNG.sequence(rngList)(rng))
  println(RNG.nonNegativeLessThan(10)(rng))
}