foo.hs

module Day8 where
    import Data.Maybe (catMaybes)
    -- from github
    elementAt :: (Eq t, Num t) => t -> [a] -> Maybe a
    elementAt 1 (x:_)  = Just x
    elementAt n (x:xs) = elementAt (n-1) xs
    elementAt _ _ = Nothing   

    -- also from github
    splitBy :: (Foldable t, Eq a) => a -> t a -> [[a]]
    splitBy delimiter = foldr f [[]] 
            where f c l@(x:xs) 
                             | c == delimiter = []:l
                             | otherwise = (c:x):xs

    parseInput :: String -> [Int]
    parseInput a = map read $ splitBy ' ' a 

    data Header = Header Int Int deriving (Show)
    data Tree = Tree Header [Tree] [Int] deriving (Show)

    readTree :: IO (Maybe Tree)
    readTree = snd . parseTree . parseInput <$> readFile "aoc/day8" 

    parseTree :: [Int] -> ([Int], Maybe Tree)
    parseTree  (n:m:tl) = (r2, Just $ Tree (Header n m) (catMaybes trees) meta)
        where 
            (r,trees) = parseTrees (tl, []) n
            (meta,r2) = splitAt m r
    parseTree i = (i, Nothing)
    

    parseTrees :: ([Int], [Maybe Tree]) -> Int -> ([Int],[Maybe Tree])
    parseTrees (ns, ts) n
        | n <= 0 = (ns, ts)
        | otherwise = parseTrees (r, ts ++ [t]) $ n-1
        where (r,t) = parseTree ns 
         
    sumMetadata :: Tree -> Int
    sumMetadata (Tree _ ts ms) = sum ms + sum (map sumMetadata ts)

    treeValue :: Tree -> Int
    treeValue (Tree _ [] xs) = sum xs
    treeValue (Tree _ ts xs) = sum [ treeValue e | m <- xs, 
                                                   e <- case elementAt m ts of 
                                                            (Just t) -> [t] 
                                                            Nothing -> []
                                                            ]