{-# OPTIONS_GHC -Wno-x-partial #-}

module Aoc where

aocMain' :: (Show p, Show a) => ([String] -> p) -> (p -> (a, c)) -> (p -> c -> a) -> IO ()
aocMain' parseFile part1 part2 = do
    contents <- readFile "input.txt"
    let input = parseFile $ lines contents
    let (part1Answer, data') = part1 input
    putStrLn $ "Part 1: " ++ show part1Answer
    let part2Answer = part2 input data'
    putStrLn $ "Part 2: " ++ show part2Answer

dup :: a -> (a, a)
dup a = (a, a)

aocMain :: (Show p, Show a) => ([String] -> p) -> (p -> a) -> (p -> a) -> IO ()
aocMain parseFile part1 part2 = aocMain' parseFile ((,()) . part1) (const' part2)

const' :: (a -> b) -> a -> c -> b
const' f a _ = f a

enumerate :: [a] -> [(Int, a)]
enumerate = zip [0..]

unimplemented :: (Show p, Show a) => p -> a
unimplemented = error . show

pair :: [a] -> (a, a)
pair as = let [a1, a2] = as in (a1, a2)

unpair :: (a, a) -> [a]
unpair (a, b) = [a, b]

triplet :: [a] -> (a, a, a)
triplet as = let [a1, a2, a3] = as in (a1, a2, a3)

single :: [a] -> a
single as = let [a] = as in a

both :: (a -> b) -> (a, a) -> (b, b)
both f (a, b) = (f a, f b)

lastSplit :: Int -> [a] -> ([a], [a])
lastSplit n as = splitAt (length as - n) as

remove :: Int -> [a] -> (a, [a])
remove n xs = let (front, a:rest) = splitAt n xs in (a, front ++ rest)

permutations :: [a] -> [[a]]
permutations [] = [[]]
permutations xs = concatMap (\n -> let (a, rest) = remove n xs in map (a:) $ permutations rest) [0..length xs-1]

adjacents :: [a] -> [(a, a)]
adjacents as = zip as $ tail as

-- starsAndBars n b returns all lists of non-negative integers of length `b+1` where the sum is `n`
starsAndBars :: Int -> Int -> [[Int]]
starsAndBars 0 b = [replicate (b+1) 0]
starsAndBars n 0 = [[n]]
starsAndBars n b = map (0:) (starsAndBars n (b-1)) ++ map (\xs -> head xs + 1 : tail xs) (starsAndBars (n-1) b)

-- Kernel

kernelMapRow' :: ([[a]] -> b) -> Maybe a -> [a] -> [a] -> [a] -> [b]
kernelMapRow' kernel x
    (a:r0@(b:c:_))
    (d:r1@(e:f:_))
    (g:r2@(h:i:_)) =
         kernel [[a,b,c],[d,e,f],[g,h,i]] : kernelMapRow' kernel x r0 r1 r2
kernelMapRow' kernel (Just x) [a, b] [c, d] [e, f] = [kernel [[a,b,x],[c,d,x],[e,f,x]]]
kernelMapRow' _ Nothing [_, _] [_, _] [_, _] = []

kernelMapRow :: ([[a]] -> b) -> Maybe a -> [a] -> [a] -> [a] -> [b]
kernelMapRow kernel (Just x) a b c = kernelMapRow' kernel (Just x) (x:a) (x:b) (x:c)
kernelMapRow kernel Nothing a b c = kernelMapRow' kernel Nothing a b c

kernelMap' :: ([[a]] -> b) -> Maybe a -> [[a]] -> [[b]]
kernelMap' kernel x (a:rest@(b:c:_)) = kernelMapRow kernel x a b c : kernelMap' kernel x rest
kernelMap' kernel (Just x) [a, b] = [kernelMapRow kernel (Just x) a b (map (const x) b)]
kernelMap' kernel Nothing [_, _] = []

kernelMap :: ([[a]] -> b) -> Maybe a -> [[a]] -> [[b]]
kernelMap kernel (Just x) as = kernelMap' kernel (Just x) ((map (const x) (head as)):as)
kernelMap kernel Nothing as = kernelMap' kernel Nothing as

-- Memoisation

type MemoF' a b a' b' = Eq a' => [(a', b')] -> a -> (b, [(a', b')])
type MemoF a b = MemoF' a b a b

memoise :: Eq a => MemoF a b -> MemoF a b
memoise f values a = case lookup a values of
    Just b -> (b, values)
    Nothing -> let (b, values') = f values a in (b, (a, b):values')

memoMap :: MemoF' a b a' b' -> MemoF' [a] [b] a' b'
memoMap f values [] = ([], values)
memoMap f values (a:as) =
    let (b, values') = f values a in
    let (bs, values'') = memoMap f values' as in
    (b:bs, values'')

unmemo :: Eq a' => MemoF' a b a' b' -> a -> b
unmemo f = fst . f []