import Control.Monad.Reader (Reader, runReader, ask, asks, local)

main :: IO ()
main = do
  putStrLn ("Naive: " ++ show (fib_naive 10))
  putStrLn ("Fast: " ++ show (fib_acc (0,1) 10))
  putStrLn ("Reader: " ++ show (runReader (fib_reader 10) (0,1)))
  putStrLn ("Reader (again): " ++ show (runReader (fib_reader_again 10) (0,1)))
  putStrLn ("FibReader: " ++ show (unFR (my_fib_reader 10) (0,1)))

-- An inefficient fib

fib_naive :: Int -> Int
fib_naive 0 = 0
fib_naive 1 = 1
fib_naive n = fib_naive (n - 1) + fib_naive (n - 2)

-- An efficient fib that uses accumulating variables

fib_acc :: (Int, Int) -> Int -> Int
fib_acc (i, j) 0 = i
fib_acc (i, j) 1 = j
fib_acc (i, j) n = fib_acc (j, i + j) (n - 1)

-- The same efficient fib using accumulating variables,
-- but made implicit with a `Reader (Int, Int)`

fib_reader :: Int -> Reader (Int, Int) Int
fib_reader n = do
  (i, j) <- ask
  case n of
    0 -> return i
    1 -> return j
    _ -> local (\(i, j) -> (j, i + j)) (fib_reader (n - 1))

-- The same fib but refactored using `asks` and a helper function

fib_reader_again :: Int -> Reader (Int, Int) Int
fib_reader_again 0 = asks fst
fib_reader_again 1 = asks snd
fib_reader_again n = local step (fib_reader_again (n - 1))
  where
    step (i, j) = (j, i + j)

-- The same fib, but refactored using a custom Monad definition instead of Reader from `mtl`

newtype FibReader a = FR { unFR :: (Int, Int) -> a }

instance Functor FibReader where
    fmap f (FR rf) = FR (\ij -> f (rf ij))

instance Applicative FibReader where
    pure a = FR (\_ij -> a)
    (FR rf) <*> (FR ra) = FR (\ij -> rf ij (ra ij))

instance Monad FibReader where
    return = pure
    (FR ra) >>= f =
      FR (\ij -> unFR (f (ra ij)) ij)

askI, askJ :: FibReader Int
askI = FR (\(i, _j) -> i)
askJ = FR snd

localStep :: FibReader a -> FibReader a
localStep m =
  FR (\(i, j) -> let f = unFR m in
                 f (j, i + j)
     )

my_fib_reader :: Int -> FibReader Int
my_fib_reader 0 = askI
my_fib_reader 1 = askJ
my_fib_reader n = localStep (my_fib_reader (n - 1))