day5.hs

module Day5 where

import Control.Arrow
import Data.List qualified as List
import Data.Map.Strict ((!?))
import Data.Map.Strict qualified as Map
import Data.Maybe
import Data.Set qualified as Set
import Data.Tuple
import Debug.Trace

main :: IO ()
main = do
  input <- readFile "day5.txt"
  let (orderingRulesText, pagesToPrintText) = second (drop 1) $ break null $ lines input
  let orderingRules = Map.fromListWith (<>) $ fmap (second Set.singleton . swap . second (drop 1) . break (== '|')) orderingRulesText
  let pagesToPrint = fmap (split ',') pagesToPrintText

  let isValid = isValid' Set.empty
        where isValid' seen (x : xs) =
                let rules = fromMaybe Set.empty $ orderingRules !? x
                in if seen `Set.isSubsetOf` rules
                     then isValid' (Set.singleton x <> seen) xs
                     else trace ("page is not valid; " <> x <> " must come after " <> show (seen Set.\\ rules)) $ False
              isValid' _ [] = True
  print $ sum $ fmap (read @Int . middleIndex) $ filter isValid pagesToPrint

  let reorder = reorder' []
        where reorder' res (x : xs) =
                let rules = Set.toList $ fromMaybe Set.empty $ orderingRules !? x
                    misplaced = res List.\\ rules
                in if null misplaced
                     then reorder' (res <> [x]) xs
                     else let (start, end) = break (`elem` misplaced) res
                           in trace ("placing " <> x <> " ahead of " <> head end) reorder' (start <> [x] <> end) xs
              reorder' res [] = res
  print $ sum $ fmap (read @Int . middleIndex) $ fmap reorder $ filter (not . isValid) pagesToPrint

  where
    middleIndex xs = xs !! (length xs `quot` 2)

    split :: Eq a => a -> [a] -> [[a]]
    split _ [] = []
    split sep str = 
        let (left, right) = break (== sep) str
        in left : split sep (drop 1 right)