Advent of code 2022 day

day5.hs

{-# LANGUAGE OverloadedStrings #-}

module DayFive
    ( partOne, partTwo
    ) where
import qualified Data.Text as T
import qualified Data.Text.IO as TI
import qualified Data.Maybe as M
import qualified Data.List as L
import qualified Control.Monad.State as S
import qualified Control.Monad as CM

data Movement = Movement { from :: Int, to :: Int, count :: Int } deriving Show
type PilesState = [String]

executeMovement :: PilesState -> Movement -> PilesState
executeMovement currentState Movement{count=0} = currentState
executeMovement currentState Movement{from=from, to=to, count=count} =
  executeMovement movedOnce Movement { from=from, to=to, count=(count - 1)}
  where movedOnce = destinationPushed
        sourcePopped = setAt currentState from newFrom
        newFrom = (init (currentState !! from))
        destinationPushed = setAt sourcePopped to $ newDestination
        newDestination = (currentState !! to) ++ [(last (currentState !! from))]

executeMovementPartTwo currentState Movement{from=from, to=to, count=count} =
  destinationPushed
  where sourcePopped = setAt currentState from newSource
        destinationPushed = setAt sourcePopped to newDestination
        removeLastN count list = take ((length list) - count) list
        getLastN count list = drop ((length list) - count) list
        oldSource = (currentState !! from)
        oldDestination = (currentState !! to)
        newSource = (removeLastN count oldSource)
        newDestination = oldDestination ++ (getLastN count oldSource)


setAt :: [a] -> Int -> a -> [a]
setAt xs i x = take i xs ++ [x] ++ drop (i + 1) xs


nthChar index string = if (length string) > index
                       then string !! index
                       else ' '

parseInstruction instructionText = Movement { count=count, from=from, to=to}
  where splitted = T.split (== ' ') instructionText
        count = readText $ splitted !! 1
        from = (readText $ splitted !! 3) - 1
        to = (readText $ splitted !! 5) - 1

lastNumberInString string = readText $ last $ filter (/= "") splitted
  where splitted = T.split (== ' ') $ string

readText = read . T.unpack

getProblemDescription :: String -> IO (PilesState, [Movement])
getProblemDescription path =
  do
    input <- TI.readFile path
    let inputLines = T.lines input
        initialStateLines = map T.unpack $ take initialStateLineCount inputLines
        initialStateLineCount = (M.fromJust $ L.elemIndex "" $ map T.unpack inputLines) - 1
        dataIdxForNthPile nth = 1 + (nth) * 4
        howManyPiles = lastNumberInString $ inputLines !! initialStateLineCount
        dataForNthPile nth = reverse $ filter (/= ' ') $ map (nthChar (dataIdxForNthPile nth)) initialStateLines
        initialState = map dataForNthPile [0..(howManyPiles - 1)]
        instructionLines = drop (initialStateLineCount + 2)  inputLines
        instructions = map parseInstruction instructionLines
    return (initialState, instructions)



partOne :: String -> IO String
partOne path = do
  (initialState, instructions) <- getProblemDescription path
  -- haha we could use 'fold' but instead use a monad why not
  let endState = S.execState applyInstructions initialState
      applyInstructions = CM.forM_ instructions $ applyInstruction
      applyInstruction instruction = S.get >>= \currentState -> S.put (executeMovement currentState instruction)
  return $ map last endState


partTwo :: String -> IO String
partTwo path = do
  (initialState, instructions) <- getProblemDescription path
  let endState = S.execState applyInstructions initialState
      applyInstructions = CM.forM_ instructions $ applyInstruction
      applyInstruction instruction = S.get >>= \currentState -> S.put (executeMovementPartTwo currentState instruction)
  return $ map last endState