module Day11 where

import Debug.Trace
import qualified Data.Map as Map
import Data.Map (Map)

parse :: String -> Map String [String]
parse = Map.fromList . map parseLine . lines
  where parseLine :: String -> (String, [String])
        parseLine s =
          let w = words s
           in (init $ head w, tail w)

pathCount :: Map String [String] -> String -> String -> Int
pathCount edges from to =
  let m = Map.fromList [(v, ct v) | v <- Map.keys edges]
      ct v = if v == to then 1 else sum $ map (m Map.!) (edges Map.! v)
   in ct from

part1 :: String -> String
part1 input =
  let edges = Map.insert "out" [] $ parse input
   in show $ pathCount edges "you" "out"


pathCount2 :: Map String [String] -> String -> String -> Int
pathCount2 edges from to =
  let neitherMemo = Map.fromList [(v, neither v) | v <- Map.keys edges]
      neither v = if v == to then 1 else sum $ map (neitherMemo Map.!) (edges Map.! v)

      withDacMemo = Map.fromList [(v, withDac v) | v <- Map.keys edges]
      withDac v = if v == to then
                    0
                  else if v == "dac" then
                    sum $ map (neitherMemo Map.!) (edges Map.! v)
                  else
                    sum $ map (withDacMemo Map.!) (edges Map.! v)
      
      withFftMemo = Map.fromList [(v, withFft v) | v <- Map.keys edges]
      withFft v = if v == to then
                    0
                  else if v == "fft" then
                    sum $ map (neitherMemo Map.!) (edges Map.! v)
                  else
                    sum $ map (withFftMemo Map.!) (edges Map.! v)
      
      withBothMemo = Map.fromList [(v, withBoth v) | v <- Map.keys edges]
      withBoth v = if v == to then
                    0
                   else if v == "dac" then
                     sum $ map (withFftMemo Map.!) (edges Map.! v)
                   else if v == "fft" then
                     sum $ map (withDacMemo Map.!) (edges Map.! v)
                   else
                     sum $ map (withBothMemo Map.!) (edges Map.! v)
   in withBoth from

part2 :: String -> String
part2 input =
  let edges = Map.insert "out" [] $ parse input
   in show $ pathCount2 edges "svr" "out"