module Main where

import Control.Monad
import qualified Data.Array as A
import Data.Function (on)
import Data.List
import qualified Data.Map.Strict as Map
import Data.Maybe
import qualified Data.Set as Set

import Input


swap :: (a, b) -> (b, a)
swap (a, b) = (b, a)

uniq :: Eq a => [a] -> [a]
uniq (x:y:zs) | x == y = uniq (x:zs)
              | otherwise = x : uniq (y:zs)
uniq l = l

parse :: String -> (String, String)
parse s = let (pre, ')' : post) = break (== ')') s
          in (pre, post)

main :: IO ()
main = do
    input <- getInput 6
    let downedges = map parse input
        upedges = map swap downedges
        nodes = uniq $ sort $ map fst $ downedges ++ upedges

    when (any ((/= 1) . length) (groupBy ((==) `on` fst) upedges)) $
        error "Input graph is not a tree"

    let parentMap = Map.fromList upedges
        node2index = Map.fromList (zip nodes [0..])
        depthArray = A.listArray (0, length nodes - 1)
                         [case Map.lookup node parentMap of
                              Nothing -> 0
                              Just par -> depthArray A.! (node2index Map.! par) + 1
                         | node <- nodes]
        checksum = sum [d | d <- A.elems depthArray] :: Int

    print checksum

    let src = parentMap Map.! "YOU"
        dest = parentMap Map.! "SAN"
        childMap = Map.fromList (map ((,) <$> fst . head <*> map snd)
                                 $ groupBy ((==) `on` fst)
                                 $ sortOn fst downedges)
        envOf n = maybe [] pure (Map.lookup n parentMap) ++ fromMaybe [] (Map.lookup n childMap)

        expand d seen border mp
            | Set.null border = mp
            | otherwise = let newnodes = Set.fromList (concatMap envOf border) Set.\\ seen
                          in expand (d + 1)
                                    (seen <> newnodes)
                                    newnodes
                                    (foldr (\n -> Map.insert n (d + 1)) mp (Set.toList newnodes))

        distMap = expand (0 :: Int) (Set.singleton src) (Set.singleton src) (Map.singleton src 0)

    print (distMap Map.! dest)