#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [pretty-simple])"
#! nix-shell -i "ghcid -c 'ghci' -T main"

{-# OPTIONS_GHC -Wall -Wincomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-unused-top-binds -Wno-unused-imports #-}
{-# OPTIONS_GHC -Wno-unused-matches #-}
{-# LANGUAGE OverloadedStrings #-}

import Debug.Trace (trace)
import Text.Pretty.Simple

import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as I

import Data.Maybe (fromMaybe)
import Data.List (iterate')

exampleInput1 :: [Int]
exampleInput1 = [0,3,6]

exampleInput2 :: [Int]
exampleInput2 = [1,3,2]
exampleInput3 :: [Int]
exampleInput3 = [3,1,2]

input :: [Int]
input = [0,3,1,6,7,5]

type Turn = Int
type Number = Int


-- Each Number is attached to the last Turn it was seen
-- The current Turn is recorded too
-- Last spoken Number is recorded too
data Game = Game { unMem :: (IntMap Turn)
                 , unTurn :: Turn
                 , unNumber :: Number
                 }
  deriving Show

initGame :: [Int] -> Game
initGame keys = Game mem ((length keys) + 1) 0
  where
    mem = I.fromList (zip keys [1..])

next :: Game -> Game
next (Game mem turn n) = fromMaybe (Game (I.insert n turn mem) (turn+1) 0) $ do
  lastTurn <- I.lookup n mem
  pure $ Game (I.insert n turn mem) (turn+1) (turn - lastTurn)

main :: IO ()
main = do
  putStrLn ":: Day 15 - Tests 1"
  print exampleInput1
  putStrLn $ "turn 04: " <> (show $ initGame exampleInput1)
  putStrLn $ "turn 05: " <> (show $ next $ initGame exampleInput1)
  putStrLn $ "turn 06: " <> (show $ next $ next $ initGame exampleInput1)
  putStrLn $ "turn 07: " <> (show $ next $ next $ next $ initGame exampleInput1)
  putStrLn $ "turn 08: " <> (show $ next $ next $ next $ next $ initGame exampleInput1)
  putStrLn $ "turn 09: " <> (show $ next $ next $ next $ next $ next $ initGame exampleInput1)
  putStrLn $ "turn 10: " <> (show $ next $ next $ next $ next $ next $ next $ initGame exampleInput1)
  print $ last $ take 2017 $ map unNumber $ iterate next $ initGame exampleInput1
  print $ last $ take 2017 $ map unNumber $ iterate next $ initGame exampleInput2
  print $ last $ take 2017 $ map unNumber $ iterate next $ initGame exampleInput3
  putStrLn ":: Day 15 - Part 1"
  print $ last $ take (2020 - 6) $ map unNumber $ iterate next $ initGame input
  putStrLn ":: Day 15 - Part 2"
  -- print $ last $ take (30000000 - 3) $ map unNumber $ iterate' next $ initGame exampleInput1
  print $ last $ take (30000000 - 6) $ map unNumber $ iterate' next $ initGame input