module Main where

import Data.Foldable (toList)
import qualified Data.IntMap.Strict as IM
import Data.IntMap.Strict (IntMap)
import Data.List (foldl')

import Input
import Util


type State = IntMap Int

next :: Int -> Int -> State -> Int
next idx num mp = maybe 0 (idx -) (IM.lookup num mp)

vaneck :: [Int] -> [Int]
vaneck values = init values ++ expand (length values) (last values) (seed (init values))
  where
    expand idx num state =
        let num' = next idx num state
        in num : expand (idx + 1) num' (IM.insert num idx state)
    seed = foldl' go mempty . zip [1..]
      where go mp (idx, num) = IM.insert num idx mp

main :: IO ()
main = do
    input <- map read . toList . splitOn (== ',') . head <$> getInput 15 :: IO [Int]
    print (vaneck input !! 2019)
    print (vaneck input !! 29999999)