import qualified Vm
import Control.Monad.Trans.State as St
import qualified Data.Map as M
import Data.Maybe (fromJust, fromMaybe)
import Data.List.Split (chunksOf)

mapS :: s -> [St.State s a] -> [s]
mapS init [] = []
mapS init (x:xs) = let s = execState x init in s : mapS s xs

main :: IO ()
main = do
    program <- Vm.readProgram
    let states = mapS Vm.initState (fmap Vm.eval program)
    let statesM = M.fromList $ zip (map Vm.step states) (map Vm.regx states)

    print statesM

    putStr "part 1: "
    let cycles = [20, 60, 100, 140, 180, 220]
    let strengths = map (\i -> i * snd (fromJust $ (i-1) `M.lookupLE` statesM)) cycles :: [Integer]
    print $ sum strengths

    putStrLn "parts 2: "
    let finalIter = fst $ fromJust $ M.lookupMax statesM
    let signals = zip [0..] $ map (\i -> snd (fromMaybe (0,1) $ (i-1) `M.lookupLE` statesM)) [1..finalIter]
    let crt = map (\(i, s) -> let p = i `mod` 40 in p >= s-1 && p <= s+1) signals
    print signals
    mapM_ putStrLn $ chunksOf 40 $ map (\b -> if b then '#' else ' ') crt