import qualified Data.Map as M
import qualified Data.Set as S
import Data.List.Split (splitOn)
import Debug.Trace (traceShowId, traceShow, trace)
import Data.Maybe (fromMaybe)

import Bfs
import Data.List (maximumBy, permutations, delete)
import Data.Function (on)

data Valve = V { flow' :: Int, neighs' :: [String] } deriving Show
type ZMap = M.Map String Valve

data NzValve = NV { flow :: Int, dists :: M.Map String Int} deriving Show
type NzMap = M.Map String NzValve

parseValve :: String -> (String, Valve)
parseValve xs = (name, V flow tunnels) where
    name = words xs !! 1
    flow = read $ last $ splitOn "=" (head $ splitOn ";" xs)
    valves = splitOn ", " $ last (splitOn "tunnels lead to valves " xs)
    singleValve = splitOn "tunnel leads to valve " xs
    tunnels = if length singleValve == 1 then valves else [last singleValve]

search' :: ZMap -> String -> String -> Int
search' scan from to = dists M.! to where
    dists = bfs f [(from,0)] M.empty
    f x = neighs' (scan M.! x)


prepass :: ZMap -> NzMap
prepass xs = M.fromList $ concatMap f $ M.toList xs where
    f (name, v@(V flow _)) = [(name, NV flow (M.fromList $ dists name)) | name == "AA" || flow > 0]
    dists name = map (\target -> (target, search' xs name target)) (filter (/= name) allValves)
    allValves = map fst . filter (\(name, V flow _) -> name == "AA" || flow > 0) $ M.toList xs




type Memo = M.Map (S.Set String, String, Int) Int
-- search :: Map -> S.Set String -> String -> Int -> Memo -> (Int, Memo)
-- search scan active location n memo 
--     | n <= 0 = (0, M.empty)
--     | otherwise = (releasedPressure + fst best, snd best) where
--         valve = scan M.! location
--         openValve = [(S.insert location active, location, n-1) | S.notMember location active && flow valve > 0]
--         releasedPressure = sum $ map (flow . (scan M.!)) (S.toList active)
--         nop = [(active, location, n-1) | S.size active == M.size scan]
--         move = map (\(name, dist) -> (active, name, n-dist)) (dists $ scan M.! location)
--         actions = openValve ++ move ++ nop
--         trace x = if n>=15 then traceShow (fst x) x else x
--         best = trace $ maximumBy (compare `on` fst) $ map (\(a,l,n) -> search scan a l n memo) actions
--         -- best = maybe best' (\b -> (b,memo)) (M.lookup (active, location, n) memo)
--         -- best' = foldr f (0, memo) actions where
--         --     f (a,l, newN) (i,m) = let (result, memo') = search scan a l newN m; b = max i result in
--         --         trace (b, M.insertWith max (a,l, newN) b $ M.unionWith max m memo')

perms :: Int -> NzMap -> [[String]]
perms limit m = perms' "AA" limit elems where
    elems = filter (/= "AA") $ M.keys m
    perms' :: String -> Int -> [String] -> [[String]]
    perms' current remaining [] = [[]]
    perms' current remaining as = do
        a <- as
        let distance = dists (m M.! current) M.! a
        if distance < remaining then do
            let l = delete a as
            ls <- perms' a (remaining-distance) l
            pure $ a:ls
        else pure [a]

execPerm :: NzMap -> [String] -> Int -> String -> Int -> Int
execPerm nzmap [] flowrate location remaining = remaining * flowrate
execPerm nzmap (target:qs) flowrate location remaining
    | remaining < 0 = 0
    | distance >= remaining = remaining * flowrate
    | otherwise = (distance * flowrate) + execPerm nzmap qs flowrate' target (remaining - distance)
    where
        valve = nzmap M.! location
        distance = 1 + dists valve M.! target
        flowrate' = flowrate + flow (nzmap M.! target)


traceMap :: (Show a, Show b) => (a -> a -> Bool) -> (a -> b) -> [a] -> [b]
traceMap cond f [] = []
traceMap cond f [x] = [f x]
traceMap cond f (x:y:xs)
    | cond x y = let result = f x in traceShow (result, x) (result: traceMap cond f (y:xs))
    | otherwise = f x : traceMap cond f (y:xs)

main :: IO ()
main = do
    x <- prepass . M.fromList . map parseValve . lines <$> getContents
    mapM_ print x

    let all = perms 30 x
    let exec p = execPerm x p 0 "AA" 30
    let result = traceMap (\a b -> take 2 a /= take 2 b) exec all
    putStr "part 1: "
    print $ maximum result
    -- print $ fst $ search x S.empty "AA" 30 M.empty