import Aoc
import Data.List.Split ( splitOn )
import Control.Arrow ( second, first )
import Data.List ( nub, partition, sort )
import Data.Maybe ( fromMaybe )

parseFile :: [String] -> ([(Int, [Int])], [[Int]])
parseFile lines =
    let 
        [rules, pages] = splitOn [""] lines
        rules' = map (pair . map read . splitOn "|") rules
        mentioned = nub $ concatMap unpair rules'
        getAll a = map snd $ filter ((==a) . fst) rules' 
    in
    (zip mentioned (map getAll mentioned), map (map read . splitOn ",") pages)
    
takeMiddle :: [a] -> a
takeMiddle xs = xs !! (length xs `div` 2)

isSafe :: [(Int, [Int])] -> [Int] -> Bool
isSafe rules = fst . foldl isSafe' (True, [])
    where
        isSafe' (False, prev) _ = (False, prev)
        isSafe' (True, prev) a = (fromMaybe True $ not . any (`elem` prev) <$> lookup a rules, a:prev)

part1 :: ([(Int, [Int])], [[Int]]) -> (Int, [[Int]])
part1 (rules, pages) = first (sum . map takeMiddle) $ partition (isSafe rules) pages

part2 :: ([(Int, [Int])], [[Int]]) -> [[Int]] -> Int
part2 (rules, _) = sum . map (takeMiddle . reorder)
    where
        getPosM rules = memoise $ \values n ->
            let ns = fromMaybe [] $ lookup n rules in
            if null ns 
                then (0, values)
                else
                    let (ns', values') = memoMap (getPosM rules) values ns in
                    (maximum ns' + 1, values')

        reorder ns =
            let rules' = map (second $ filter (`elem` ns)) $ filter ((`elem` ns) . fst) rules in
            map snd $ sort $ zip (unmemo (memoMap $ getPosM rules') ns) ns

main :: IO ()
main = aocMain' parseFile part1 part2