AOC.hs

{-|
Module      : AOC
Description : Utility functions commonly used while solving Advent of Code puzzles
Copyright   : (c) M1n3c4rt, 2025
License     : BSD-3-Clause
Maintainer  : vedicbits@gmail.com
Stability   : stable
-}

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TupleSections #-}
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-type-defaults #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveTraversable #-}

module Utility.AOC (
    -- * Pathfinding algorithms
    -- $cat1
    shortestDistance,
    shortestPaths,
    shortestDistanceWith,
    shortestPathsWith,
    -- * Neighbour functions
    neighbours4,
    neighbours8,
    neighbours6,
    neighbours26,
    -- * Taxicab (Manhattan) distance
    taxicab2,
    taxicab3,
    taxicab,
    -- * Input parsing
    -- $cat2
    enumerate',
    enumerateRead',
    enumerate,
    enumerateRead,
    enumerateHM,
    enumerateReadHM,
    enumerateFilter,
    enumerateFilterSet,
    numbers,
    numbers',
    -- * Flood fill
    floodFill,
    floodFillWith,
    -- * List selection
    choose,
    permute,
    takeEvery,
    chunk,
    -- * Extrapolation
    extrapolate,
    -- * Debugging
    prettyPrintSet,
    prettyPrintSetWide,
    prettyPrintHM,
    prettyPrintHMWide,
    traceSleep,
    traceSleepSeconds,
    -- * Memoization
    -- $cat3
    memo2,
    memo3,
    memo4,
    -- * ComplexL1 datatype
    ComplexL1 (..),
    realPart,
    imagPart,
    conjugate,
    magnitude,
    -- * Miscellaneous
    range,
    rangeIntersect,
    binToDec,
    -- Export types/constraints used in top-level function type signatures
    Hashable,
    HM.HashMap,
    S.Set
) where

import qualified Data.HashMap.Strict as HM
import Data.Hashable (Hashable (hashWithSalt))
import qualified Data.Set as S
import qualified Data.Heap as H
import Data.List (permutations, genericIndex, groupBy)
import Data.Maybe (fromMaybe)
import System.IO.Unsafe (unsafePerformIO)
import Control.Concurrent (threadDelay, newMVar, readMVar, modifyMVar_)
import Data.Function (on)
import Data.Data (Data)
import GHC.Generics (Generic, Generic1)
import Foreign (Storable (sizeOf, peek, poke, peekElemOff, pokeElemOff, alignment), castPtr)
import qualified GHC.Ptr
import Control.Monad.Zip (MonadZip (mzipWith))
import Control.Monad.Fix (MonadFix (mfix))

createMinPrioHeap :: Ord a1 => (a1,a) -> H.MinPrioHeap a1 a
createMinPrioHeap = H.singleton

-- $cat1
-- All of the following functions return distances as a @Maybe Int@, where @Nothing@ is returned if no path is found.
-- The graph is a @HashMap@ mapping each node to a sequence of (neighbour, edge weight) pairs.
-- Functions that take a graph assume that any neighbour that a node in the graph points to is also in the graph.

-- | Returns the shortest distance between two nodes in a graph.
shortestDistance :: (Foldable t, Hashable n, Ord a, Num a)
    => HM.HashMap n (t (n, a)) -- ^ Graph
    -> n -- ^ Start node
    -> n -- ^ End node
    -> Maybe a
shortestDistance graph = shortestDistanceWith (\n -> fromMaybe (error "Node not in graph") $ HM.lookup n graph)

-- | Returns the shortest distance between two nodes in a graph and a list of all possible paths from the ending node to the starting node.
-- The starting node is not included in each path.
shortestPaths :: (Foldable t, Hashable n, Ord a, Num a)
    => HM.HashMap n (t (n, a)) -- ^ Graph
    -> n -- ^ Start node
    -> n -- ^ End node
    -> (Maybe a, [[n]])
shortestPaths graph = shortestPathsWith (\n -> fromMaybe (error "Node not in graph") $ HM.lookup n graph)

-- | Given a function that takes a node and returns a sequence of (neighbour,edge weight) pairs, returns the shortest distance between two nodes.
shortestDistanceWith :: (Foldable t, Hashable n, Ord a, Num a)
    => (n -> t (n, a)) -- ^ Function to generate neighbours
    -> n -- ^ Start node
    -> n -- ^ End node
    -> Maybe a
shortestDistanceWith f start end = fst $ shortestPathsWith f start end

-- | Given a function that takes a node and returns a sequence of (neighbour,edge weight) pairs, returns the shortest distance between two nodes and a list of all possible paths from the ending node to the starting node.
-- The starting node is not included in each path.
shortestPathsWith :: (Foldable t, Hashable n, Ord a, Num a)
    => (n -> t (n, a)) -- ^ Function to generate neighbours
    -> n -- ^ Start node
    -> n -- ^ End node
    -> (Maybe a, [[n]])
shortestPathsWith f start end =
    let initQueue = createMinPrioHeap (0,start)
        initPaths = HM.singleton start (0,[[]])
        helper (paths,queue) = case H.view queue of
            Nothing -> (paths,queue)
            Just ((_,n),ns) ->
                let Just (currentDistance,currentPaths) = HM.lookup n paths
                    neighbours = f n
                    updateNeighbour (n',d') (p',q') = case HM.lookup n' p' of
                        Nothing -> (HM.insert n' (currentDistance+d',map (n':) currentPaths) p', H.insert (currentDistance+d',n') $ H.filter ((/=n') . snd) q')
                        Just (d'',ps'') ->
                            if d'' < currentDistance+d' then
                                (p',q')
                            else if d'' > currentDistance+d' then
                                (HM.insert n' (currentDistance+d',map (n':) currentPaths) p', H.insert (currentDistance+d',n') $ H.filter ((/=n') . snd) q')
                            else
                                (HM.insert n' (currentDistance+d',ps'' ++ map (n':) currentPaths) p', q')
                in helper $ foldr updateNeighbour (paths,ns) neighbours

    in case HM.lookup end $ fst (helper (initPaths,initQueue)) of
        Nothing -> (Nothing, [])
        Just (d,ps) -> (Just d, ps)

-- | Returns the 4 points orthogonally adjacent to the given point.
neighbours4 :: (Num a, Num b) => (a, b) -> [(a, b)]
neighbours4 (x,y) = [(x+1,y),(x,y+1),(x-1,y),(x,y-1)]

-- | Returns the 8 points orthogonally or diagonally adjacent to the given point.
neighbours8 :: (Eq a, Eq b, Num a, Num b) => (a, b) -> [(a, b)]
neighbours8 (x,y) = [(x+p,y+q) | p <- [-1,0,1], q <- [-1,0,1], p /= 0 || q /= 0]

-- | Returns the 6 points orthogonally adjacent to the given point in 3D space.
neighbours6 :: (Num a, Num b, Num c) => (a, b, c) -> [(a, b, c)]
neighbours6 (x,y,z) = [(x+1,y,z),(x,y+1,z),(x,y,z+1),(x-1,y,z),(x,y-1,z),(x,y,z-1)]

-- | Returns the 26 points orthogonally or diagonally adjacent to the given point in 3D space.
neighbours26 :: (Eq a, Eq b, Eq c, Num a, Num b, Num c) => (a, b, c) -> [(a, b, c)]
neighbours26 (x,y,z) = [(x+p,y+q,z+r) | p <- [-1,0,1], q <- [-1,0,1], r <- [-1,0,1], p /= 0 || q /= 0 || r /= 0]

-- | Returns the Taxicab/Manhattan distance between two points in 2D space.
taxicab2 :: Num a => (a, a) -> (a, a) -> a
taxicab2 (a,b) (c,d) = abs (a-c) + abs (b-d)

-- | Returns the Taxicab/Manhattan distance between two points in 3D space.
taxicab3 :: Num a => (a, a, a) -> (a, a, a) -> a
taxicab3 (a,b,c) (d,e,f) = abs (a-d) + abs (b-e) + abs (c-f)

-- | Returns the Taxicab/Manhattan distance between two points in n dimensions, where both points are lists of length n.
taxicab :: Num a => [a] -> [a] -> a
taxicab as bs = sum $ zipWith (\x y -> abs (x-y)) as bs

-- $cat2
-- The following functions (beginning with "enumerate") operate on a grid of characters as a string with a newline after each row (as seen in several Advent of Code puzzle inputs).

-- | Converts a grid to a list of pairs @((x,y),c)@ representing xy coordinates and the character at that location.
enumerate' :: (Num y, Num x) => String -> [((x, y), Char)]
enumerate' s =
    let ss = lines s
        ys = zipWith (\n l -> map (n,) l) (iterate (+1) 0) ss
        xs = map (zipWith (\x (y,c) -> ((x,y),c)) (iterate (+1) 0)) ys
    in concat xs

-- | Enumerates a grid along with reading the characters (usually as integers), and returns a list of pairs.
enumerateRead' :: (Read c, Num y, Num x) => String -> [((x, y), c)]
enumerateRead' = map (\((x,y),c) -> ((x,y),read [c])) . enumerate'

-- | Converts a grid to a list of triples @(x,y,c)@ representing xy coordinates and the character at that location.
enumerate :: (Num y, Num x) => String -> [(x, y, Char)]
enumerate = map (\((x,y),c) -> (x,y,c)) . enumerate'

-- | Enumerates a grid along with reading the characters (usually as integers), and returns a list of triples.
enumerateRead :: (Read c, Num y, Num x) => String -> [(x, y, c)]
enumerateRead = map (\((x,y),c) -> (x,y,read [c])) . enumerate'

-- | Enumerates a grid and stores it in a @HashMap@ where points are mapped to the character at that location.
enumerateHM :: (Num x, Num y, Enum x, Enum y, Hashable x, Hashable y) => String -> HM.HashMap (x, y) Char
enumerateHM = HM.fromList . enumerate'

-- | Enumerates a grid and stores it in a @HashMap@ along with reading the characters (usually as integers).
enumerateReadHM :: (Num x, Num y, Enum x, Enum y, Hashable x, Hashable y, Read c) => String -> HM.HashMap (x, y) c
enumerateReadHM = HM.fromList . map (\((x,y),c) -> ((x,y),read [c])) . enumerate'

-- | Returns a list of points on a grid for which a certain condition is met.
enumerateFilter :: (Num y, Num x) => (Char -> Bool) -> String -> [(x, y)]
enumerateFilter f = map fst . filter (f . snd) . enumerate'

-- | Returns a set of points on a grid for which a certain condition is met.
enumerateFilterSet :: (Ord x, Ord y, Num y, Num x) => (Char -> Bool) -> String -> S.Set (x, y)
enumerateFilterSet f = S.fromList . enumerateFilter f

-- | Returns all the integers in a string (including negative signs).
numbers :: (Num a, Read a) => [Char] -> [a]
numbers = map read . filter (isDigit . head) . groupBy ((==) `on` isDigit)
    where isDigit = (`elem` "1234567890-")

-- | Returns all the integers in a string (excluding negative signs).
numbers' :: (Num a, Read a) => [Char] -> [a]
numbers' = map read . filter (isDigit . head) . groupBy ((==) `on` isDigit)
    where isDigit = (`elem` "1234567890")

floodFill' :: Ord a => (a -> [a]) -> S.Set a -> [a] -> S.Set a -> S.Set a
floodFill' neighbours finished (f:frontier) blocks = floodFill' neighbours (S.insert f finished) (frontier++filtered) blocks
    where filtered = filter (\n -> n `S.notMember` finished && n `notElem` frontier && n `S.notMember` blocks) $ neighbours f
floodFill' _ finished [] _ = finished

floodFillWith' :: Ord a => (a -> a -> Bool) -> (a -> [a]) -> S.Set a -> [a] -> S.Set a
floodFillWith' cond neighbours finished (f:frontier) = floodFillWith' cond neighbours (S.insert f finished) (frontier++filtered)
    where filtered = filter (\n -> n `S.notMember` finished && n `notElem` frontier && cond f n) $ neighbours f
floodFillWith' _ _ finished [] = finished

-- | Applies a flood fill algorithm given a function to generate a point's neighbours, a starting set of points, and a set of points to avoid. Returns a set of all points covered.
floodFill :: Ord a
    => (a -> [a]) -- ^ Neighbour function
    -> S.Set a -- ^ Initial set of points
    -> S.Set a -- ^ Set of points to avoid
    -> S.Set a
floodFill neighbours frontier = floodFill' neighbours S.empty (S.toList frontier)

-- | Applies a flood fill algorithm given a function to generate a point's neighbours, a condition that filters out points generated by said function, and a starting set of points. Returns a set of all points covered.
-- The condition is of the form @a -> a -> Bool@, which returns @True@ if the second point is a valid neighbour of the first point and @False@ otherwise.
floodFillWith :: Ord a
    => (a -> a -> Bool) -- ^ Condition
    -> (a -> [a]) -- ^ Neighbour function
    -> S.Set a -- ^ Initial set of points
    -> S.Set a
floodFillWith cond neighbours frontier = floodFillWith' cond neighbours S.empty (S.toList frontier)

-- | Generates a list of all possible lists of length n by taking elements from the provided list of length l.
-- Relative order is maintained, and the length of the returned list is \(_{n}C_{l}\).
choose :: (Num n, Ord n) => n -> [a] -> [[a]]
choose 0 _ = [[]]
choose _ [] = []
choose n (x:xs)
    | n > fromIntegral (length (x:xs)) = []
    | otherwise = map (x:) (choose (n-1) xs) ++ choose n xs

-- | Generates a list of all possible lists of length n by taking elements from the provided list of length l.
-- The length of the returned list is \(_{n}P_{l}\).
permute :: (Num n, Ord n) => n -> [a] -> [[a]]
permute n = concatMap permutations . choose n

-- | Takes every nth element from a list xs, starting from @xs !! (n-1)@.
takeEvery :: Int -> [a] -> [a]
takeEvery _ [] = []
takeEvery n xs = let (a,b) = splitAt n xs in if length a < n then [] else last a:takeEvery n b

-- | Splits a list into sublists of size n. The length of the last sublist may be less than n.
chunk :: Int -> [a] -> [[a]]
chunk _ [] = []
chunk n xs = let (a,b) = splitAt n xs in a:chunk n b

-- | Gets the nth element of an infinite list, assuming that each element in the list can be generated using the previous element, for example, a list generated with @iterate@.
extrapolate :: (Integral b, Ord a) => b -> [a] -> a
extrapolate n ls = let (o,p) = helper 0 S.empty ls in ls `genericIndex` (((n-o) `mod` p) + o)
    where
        helper k finished (l:ls')
            | S.null matches = helper (k+1) (S.insert (k,l) finished) ls'
            | otherwise = let o = fst $ S.elemAt 0 matches in (o,k-o)
            where matches = S.filter ((==l) . snd) finished

-- | Converts a set of points @(x,y)@ to a string composed of @'#'@ and @' '@. This function is useful when displaying puzzle answers formed by a grid of points.
-- Up to translation of points, @prettyPrintSet . enumerateFilterSet (=='#') = id@.
prettyPrintSet :: (Enum b, Enum a, Ord a, Ord b) => S.Set (a, b) -> String
prettyPrintSet points = unlines [[if (x,y) `S.member` points then '#' else ' ' | x <- [xmin..xmax]] | y <- reverse [ymin..ymax]]
    where
        xs = S.map fst points
        ys = S.map snd points
        (xmin,xmax,ymin,ymax) = (minimum xs, maximum xs, minimum ys, maximum ys)

-- | Same as @prettyPrintSet@, but displays points at double width to improve readability.
prettyPrintSetWide :: (Enum b, Enum a, Ord a, Ord b) => S.Set (a, b) -> String
prettyPrintSetWide = foldr (\c acc -> if c /= '\n' then c:c:acc else c:acc) [] . prettyPrintSet

-- | Converts a @HashMap@ of points @(x,y)@ and characters @c@ to a string with the corresponding character at each point. This function is useful when displaying puzzle answers formed by a grid of points.
-- Up to translation of points, @prettyPrintHM . enumerateHM = id@.
prettyPrintHM :: (Enum b, Enum a, Hashable a, Hashable b, Ord a, Ord b) => HM.HashMap (a, b) Char -> String
prettyPrintHM points = unlines [[HM.lookupDefault ' ' (x,y) points | x <- [xmin..xmax]] | y <- reverse [ymin..ymax]]
    where
        xs = map fst $ HM.keys points
        ys = map snd $ HM.keys points
        (xmin,xmax,ymin,ymax) = (minimum xs, maximum xs, minimum ys, maximum ys)

-- | Same as @prettyPrintHM@, but displays points at double width to improve readability.
prettyPrintHMWide :: (Enum b, Enum a, Hashable a, Hashable b, Ord a, Ord b) => HM.HashMap (a, b) Char -> String
prettyPrintHMWide = foldr (\c acc -> if c /= '\n' then c:c:acc else c:acc) [] . prettyPrintHM

{-# NOINLINE traceSleep #-}
-- | Pauses execution for n microseconds, before returning the second argument as its result. Useful for slowing down output that normally floods the terminal.
-- Like functions exported by Debug.Trace, this function should only be used for debugging.
-- The function is not referentially transparent: its type indicates that it is a pure function but it has the side effect of delaying execution.
traceSleep :: Int -> a -> a
traceSleep n x = unsafePerformIO $ do
    threadDelay n
    return x

{-# NOINLINE traceSleepSeconds #-}
-- | Pauses execution for n seconds. See @traceSleep@.
traceSleepSeconds :: Int -> a -> a
traceSleepSeconds n = traceSleep (n*1000000)

-- $cat3
-- Memoize a function with multiple arguments. Uses @memo@ from @Data.MemoUgly@ with slight modifications.

memo2 :: (Hashable a, Hashable b) => (a -> b -> c) -> (a -> b -> c)
memo2 f = unsafePerformIO $ do
    v <- newMVar HM.empty
    let f' a b = unsafePerformIO $ do
            m <- readMVar v
            case HM.lookup (a,b) m of
                Nothing -> do let { r = f a b }; modifyMVar_ v (return . HM.insert (a,b) r); return r
                Just r  -> return r
    return f'
memo3 :: (Hashable a, Hashable b, Hashable c) => (a -> b -> c -> d) -> (a -> b -> c -> d)
memo3 f = unsafePerformIO $ do
    v <- newMVar HM.empty
    let f' a b c = unsafePerformIO $ do
            m <- readMVar v
            case HM.lookup (a,b,c) m of
                Nothing -> do let { r = f a b c }; modifyMVar_ v (return . HM.insert (a,b,c) r); return r
                Just r  -> return r
    return f'
memo4 :: (Hashable a, Hashable b, Hashable c, Hashable d) => (a -> b -> c -> d -> e) -> (a -> b -> c -> d -> e)
memo4 f = unsafePerformIO $ do
    v <- newMVar HM.empty
    let f' a b c d = unsafePerformIO $ do
            m <- readMVar v
            case HM.lookup (a,b,c,d) m of
                Nothing -> do let { r = f a b c d }; modifyMVar_ v (return . HM.insert (a,b,c,d) r); return r
                Just r  -> return r
    return f'

-- | Generates a range with @[x..y]@, but reverses the list instead of returning an empty range if x > y.
range :: (Ord a, Enum a) => a -> a -> [a]
range x y = if y < x then [x,pred x..y] else [x..y]

-- | Takes (a,b) and (c,d) as arguments and returns the intersection of the ranges [a..b] and [c..d] as another pair if it is not empty.
rangeIntersect :: Ord b => (b, b) -> (b, b) -> Maybe (b, b)
rangeIntersect (a,b) (c,d)
    | b < c || a > d = Nothing
    | otherwise = Just (max a c, min b d)

-- | Converts a list of booleans (parsed as a binary number) to an integer.
binToDec :: Num a => [Bool] -> a
binToDec = sum . zipWith (*) (map (2^) [0..]) . map (fromIntegral . fromEnum) . reverse

infix 6 :+
-- | A @Complex@ number type whose instance requires the constituent type to only be part of @Num@ instead of @RealFloat@.
-- As a consequence of this, the @abs@ and @magnitude@ functions use the L1 metric instead of the L2 metric (which is quite useful for Advent of Code puzzles).
-- Also, @signum@ is performed element-wise and can only return one of 8 different values.
-- This breaks the equality @abs z * signum z == z@ that is true in the Complex datatype.
-- The functions @mkPolar@, @cis@, @polar@ and @phase@ are unimplemented for obvious reasons.
-- Unlike @Complex@, Ord and Hashable are supported. The Ord instance treats the values as tuples.
-- The name of the datatype contains "L1" in order to reflect these changes.
data ComplexL1 a = a :+ a deriving (Eq, Ord, Read, Show, Data, Generic, Generic1, Functor, Foldable, Traversable)

-- | Real part of a complex number.
realPart :: ComplexL1 a -> a
realPart (a :+ _) = a
-- | Imaginary part of a complex number.
imagPart :: ComplexL1 a -> a
imagPart (_ :+ b) = b

-- | Conjugate of a complex number.
conjugate :: Num a => ComplexL1 a -> ComplexL1 a
conjugate (a :+ b) = a :+ negate b
-- | (Taxicab) magnitude of a complex number.
magnitude :: Num a => ComplexL1 a -> a
magnitude (a :+ b) = abs a + abs b

instance Num a => Num (ComplexL1 a) where
    (+) :: Num a => ComplexL1 a -> ComplexL1 a -> ComplexL1 a
    (a:+b) + (c:+d) = (a+c) :+ (b+d)
    (*) :: Num a => ComplexL1 a -> ComplexL1 a -> ComplexL1 a
    (a:+b) * (c:+d) = (a*c - b*d) :+ (b*c + a*d)
    abs :: Num a => ComplexL1 a -> ComplexL1 a
    abs (a :+ b) = (abs a + abs b) :+ 0
    signum :: Num a => ComplexL1 a -> ComplexL1 a
    signum (a :+ b) = signum a :+ signum b
    fromInteger :: Num a => Integer -> ComplexL1 a
    fromInteger n = fromInteger n :+ 0
    negate :: Num a => ComplexL1 a -> ComplexL1 a
    negate (a :+ b) = (a :+ b) * ((-1) :+ 0)

instance Hashable a => Hashable (ComplexL1 a) where
    hashWithSalt :: Hashable a => Int -> ComplexL1 a -> Int
    hashWithSalt salt (a :+ b) = hashWithSalt (hashWithSalt salt a) b

instance Storable a => Storable (ComplexL1 a) where
    sizeOf :: Storable a => ComplexL1 a -> Int
    sizeOf a       = 2 * sizeOf (realPart a)
    alignment :: Storable a => ComplexL1 a -> Int
    alignment a    = alignment (realPart a)
    peek :: Storable a => GHC.Ptr.Ptr (ComplexL1 a) -> IO (ComplexL1 a)
    peek p           = do
                        let q = castPtr p
                        r <- peek q
                        i <- peekElemOff q 1
                        return (r :+ i)
    poke :: Storable a => GHC.Ptr.Ptr (ComplexL1 a) -> ComplexL1 a -> IO ()
    poke p (r :+ i)  = do
                        let q = castPtr p
                        poke q r
                        pokeElemOff q 1 i

instance Applicative ComplexL1 where
  pure :: a -> ComplexL1 a
  pure a = a :+ a
  (<*>) :: ComplexL1 (a -> b) -> ComplexL1 a -> ComplexL1 b
  f :+ g <*> a :+ b = f a :+ g b
  liftA2 :: (a -> b -> c) -> ComplexL1 a -> ComplexL1 b -> ComplexL1 c
  liftA2 f (x :+ y) (a :+ b) = f x a :+ f y b

instance Monad ComplexL1 where
  a :+ b >>= f = realPart (f a) :+ imagPart (f b)

instance MonadZip ComplexL1 where
  mzipWith = liftA2

instance MonadFix ComplexL1 where
  mfix f = (let a :+ _ = f a in a) :+ (let _ :+ a = f a in a)