import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import Debug.Trace
import Control.Monad.State

-- 1st is turn number
-- 2nd is map of spoken
-- 3rd is next number to speak
type S = (Int, Map Int Int, Int)

produceInitial lst =
    let len = length lst
        f acc (idx, el) = Map.insert el idx acc
        initMap = foldl f Map.empty $ zip [0..] lst
    in (len, initMap, 0)

speak :: Int -> State S Int
speak n = do
    (turn, spoken, nextSpeak) <- get
    let lastSpoke = Map.lookup nextSpeak spoken
    let spoken2 = Map.insert nextSpeak turn spoken
    let nextSpeak2 = case lastSpoke of
            Nothing -> 0
            Just n -> turn - n
    let result = (turn + 1, spoken2, nextSpeak2)
    -- modify (\_ -> traceShow result result)
    modify (\_ -> result)
    return nextSpeak
    
speakStream lst = (++) lst $ flip evalState initial $ mapM speak [0..]
    where initial = produceInitial lst

actualInput = [5,1,9,18,13,8,0]
exInput = [0,3,6]
n = 30000000
-- n = 2020

main = do print (speakStream actualInput !! (n-1))