Day6.hs

{-# language DeriveFunctor #-}
{-# language TupleSections #-}

module Day6 where
  
import Data.List as L
import Data.List.Split
import Data.Maybe
import Data.Map as M

-- The whole recursion-schemes library I need

newtype Fix f = Fix { unFix :: f (Fix f) }
type Coalgebra f a = a -> f a
ana :: Functor f => Coalgebra f a -> a -> Fix f
ana coalg = Fix . fmap (ana coalg) . coalg
type Algebra f a = f a -> a
cata :: Functor f => Algebra f a -> Fix f -> a
cata alg = alg . fmap (cata alg) . unFix
hylo :: Functor f => Algebra f a -> Coalgebra f b -> b -> a
hylo f g = f . fmap (hylo f g) . g

-- Functor that generates the rose tree

data TreeF a = NodeF String [a]
  deriving (Functor, Show)
  
-- Rose tree
type Tree = Fix TreeF

-- A mulitmap: Who orbits a given name (directly) 
type Pool = M.Map String [String]

-- Unfolding the tree top down

mkNode :: Coalgebra TreeF (String, Pool)
mkNode (key, pool) = 
  case M.lookup key pool of
    Nothing  -> NodeF key []  -- Leaf
    Just lst -> -- list of children
      let newPool = M.delete key pool
      in NodeF key (fmap ( , newPool) lst) -- tuple section

-- Folding the tree bottom up

-- The accumulator accumulates (#of children, # of orbits)

orbits :: Algebra TreeF (Int, Int) 
orbits (NodeF _ lst) = 
    (childCount, orbitCount + childCount)
  where childCount = sum $ fmap ((+1) . fst) lst
        orbitCount = sum $ fmap snd lst
 

allOrbits = snd . hylo orbits mkNode . ("COM", )

-- Part II

-- Folding the tree bottom up

-- The accumulator
data Accum = Acc { seenMe :: Bool
                 , seenSa :: Bool
                 , dist  :: Int } 
  deriving Show

-- Combine info gathered from children

meAndSanta :: Algebra TreeF Accum
meAndSanta (NodeF name accums) 
  | name == "YOU" || name == "SAN"
  = if name == "YOU"
    then if hasSanta accums
         then Acc True  True (distToSa accums)
         else Acc True  False 0
    else if hasMe accums
         then Acc True  True (distToMe accums)
         else Acc False True 0      
  where hasMe    = or . fmap seenMe
        hasSanta = or . fmap seenSa
        distToMe = dist . fromJust . find seenMe
        distToSa = dist . fromJust . find seenSa
        
meAndSanta (NodeF name accums) = -- neither YOU nor SANTA
  case (ixMe, ixSa) of
     (Nothing, Nothing)     -> Acc False False (-1)
     (Nothing, Just idx2)   -> Acc False True  (1 + dist (accums!! idx2))
     (Just idx1, Nothing)   -> Acc True  False (1 + dist (accums!! idx1))
     (Just idx1, Just idx2) -> -- Both have been seen!
         if idx1 == idx2  -- in the same child tree
         then Acc True True ( dist (accums!! idx1))
         else Acc True True ((dist (accums!! idx1)) + (dist (accums!! idx2)))
  where ixMe = L.findIndex seenMe accums
        ixSa = L.findIndex seenSa accums

meetSanta = dist . hylo meAndSanta mkNode . ("COM", )

-- for testing only
allOrbitsFromTree = cata orbits
meetingFromTree = cata meAndSanta

-- Show tree for debugging
shTree :: Tree -> String
shTree = cata showNode 
  where
    showNode :: Algebra TreeF String
    showNode (NodeF s lst) = s ++ "->[" ++ intercalate "," lst ++ "]"
    
-- A handy multimap function
appendMulti :: (String, String) -> Pool -> Pool
appendMulti (key, value) pool =
  case M.lookup key pool of
    Nothing -> M.insert key [value] pool
    Just lst -> M.insert key (value:lst) pool

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

main = do
  text <- readFile "Data6.txt"
  --let text = unlines $ test2
  let pairs = fmap (toPair . splitOn ")") $ lines text
  --print pairs
  let pool = Prelude.foldl (\m kv -> appendMulti kv m) M.empty pairs
  --print pool
  --let tree = ana mkNode ("COM", pool)
  --putStrLn $ shTree tree
  print $ allOrbits pool
  print $ meetSanta   pool
  
test' = ["COM)B"]
test'' = ["COM)B", "B)C", "B)D"]
test = ["COM)B"
       , "B)C"
       , "C)D"
       , "D)E"
       , "E)F"
       , "B)G"
       , "G)H"
       , "D)I"
       , "E)J"
       , "J)K"
       , "K)L"]
       
test2 = ["COM)B"
        ,"B)C"
        ,"C)D"
        ,"D)E"
        ,"E)F"
        ,"B)G"
        ,"G)H"
        ,"D)I"
        ,"E)J"
        ,"J)K"
        ,"K)L"
        ,"K)YOU"
        ,"I)SAN"]