State Monad - RNG

-- 2024-09-06 - Emilio Francesquini <[email protected]>
-- Universidade Federal do ABC
-- CC-BY-SA 4.0

module Main where

import Data.Bits 
import Control.Monad 

-----------------------------------
-- Gerador de números aleatórios --
-----------------------------------

type RandomState = (Int, Int)

simpleRandomInit :: Int -> RandomState
simpleRandomInit seed =
  (z, w)
  where
    n1 = (seed * 48947) `mod` 4294967296;
    z = if n1 /= 9 then n1 else 362436069
    n2 = (seed * 104623) `mod` 4294967296;
    w = if n2 /= 0 then n2 else 521288629


-- Este simples gerador de números aleatórios é baseado naquele
-- proposto por George Marsaglia (MWC - multiply with carry). Apesar
-- de ser muito simples, ele é capaz de passar pela série de testes
-- Marsaglia's DIEHARD para geradores de números aleatórios. A
-- implementação abaixo é uma adaptação daquela feita por John
-- D. Cook.
nextInt :: RandomState -> (Int, RandomState)
nextInt (z, w) =
  (u, (z', w'))
  where
    z' = 36969 * (z .&. 65535) + (z `shiftR` 16)
    w' = 18000 * (w .&. 65535) + (w `shiftR` 16);
    u  = (z `shiftL` 16) + w;


---------------------------
-- Utilização aleatórios --
---------------------------

-- Usa o gerador como função pura
randomList :: Int -> [Int]
randomList seed =
  let
    s0 = simpleRandomInit seed
    (v0, s1) = nextInt s0 -- precisa carregar o estado
    (v1, s2) = nextInt s1 -- entre cada uma das chamadas
    (v2, _) = nextInt s2 in
  [v0, v1, v2]

randomListSize :: Int -> Int -> [Int]
randomListSize seed n =
  -- Para carregar o estado para um número arbitrário de chamadas usa
  -- um fold. Poderia ser também uma versão recursiva.
  reverse . fst $ foldl
    (\(res, s) _ -> let (v, s') = nextInt s in (v:res, s'))
    ([], simpleRandomInit seed)
    [1..n]


-- Transformador de estados
-- dado um estado inicial, devolve um novo estado e um valor
newtype ST s a = ST (s -> (a, s))

-- Roda o state transformer com o estado (ambos recebidos por
-- parâmetro) e devolve uma tupla com o resultado da execução e o novo
-- estado
rodaCom :: ST s a -> s -> (a, s)
rodaCom (ST f) s0 = f s0

-- Veja slides da aula para maiores explicações
instance Functor (ST s) where
  -- fmap :: (a -> b) -> ST a -> ST b
  -- x :: a
  fmap g st = ST stb
    where
      stb s = (g x, s')
        where
          (x, s') = rodaCom st s

-- Veja slides da aula para maiores explicações
instance Applicative (ST s) where
  -- <*> :: ST (a -> b) -> ST a -> ST b
  stf <*> stx = ST stb
    where
      stb s = (f x, s'')
        where
          (f, s')  = rodaCom stf s
          (x, s'') = rodaCom stx s'

  pure x = ST (\s -> (x, s))

-- Veja slides da aula para maiores explicações
instance Monad (ST s) where
  -- (>>=) :: ST a -> (a -> ST b) -> ST b
  st >>= f = ST stb
    where
      stb s = rodaCom (f x) s'
        where (x, s') = rodaCom st s



-- Quero uma função que dado um estado me devolva um aleatório +
-- estado alterado

nextIntM :: ST RandomState Int
nextIntM = ST nextInt

-- Devolve um transformador de estados ST [Int] que quando executado
-- com a função rodaCom devolverá ([Int], State) Para ser justo (e
-- comparável com a função randomList) não usa mapM. A próxima função
-- corrige isto.
randomListM' :: ST RandomState [Int]
randomListM' = do
    v0 <- nextIntM
    v1 <- nextIntM
    v2 <- nextIntM
    return [v0, v1, v2]

randomListM :: Int -> [Int]
randomListM seed =
  -- - rodaCom roda o transformador de estados randomList2' tendo como
  --   estado inicial o resultado de (simpleRandomInit seed)
  -- - fst "tira" o resultado do (a, State) = Valor + contexto
  fst $ rodaCom randomListM'(simpleRandomInit seed)

-- Dado n, a quantidade de números aleatórios a serem gerados, devolve
-- um transformador de estados que quando executado com a função
-- rodaCom devolve um ([Int], State)
randomListSizeM' :: Int -> ST RandomState [Int]
randomListSizeM' n = replicateM n nextIntM

-- Código praticamente idêntico ao randomList2
randomListSizeM :: Int -> Int -> [Int]
randomListSizeM seed n =
  fst $ rodaCom (randomListSizeM' n) (simpleRandomInit seed)


-- Estado é dado por (fibx, (penultimo,ultimo))

-- Um passo de transformação
nextFiboM :: ST (Integer, Integer) Integer
nextFiboM = ST fib
  where
    fib (a, b) = (a + b, (b, a + b))

-- >>> rodaCom nextFiboM (2,3)
-- (5,(3,5))

-- Vários passos
fibosM :: ST (Integer, Integer) [Integer]
fibosM = sequence $ repeat nextFiboM

fibos :: [Integer]
fibos = 0 : 1 : fst (rodaCom fibosM (0, 1))

-- >>> take 100 $ fibos
-- [0,1,1,2,3,5,8,13,21,34]

main :: IO ()
main = do
  putStrLn "Aleatórios"
  print $ randomList 42
  print $ randomListSize 42 3
  print $ randomListM 42
  print $ randomListSizeM 42 3