module Commons where

import Data.List ((\\), delete, nub)
import Data.List.Utils (split)
import Data.Map (Map, singleton, insert, filterWithKey, empty, unionWith, fromList, union, (!), keys, member, size)
import GHC.IO.Handle (isEOF)


type Component = String
type Network = Map Component [Component]


parseLine :: Component -> [Component] -> Network
parseLine c dest = foldl (\ n newC -> insert newC [c] n) (singleton c dest) dest

parse :: IO Network
parse = do done <- isEOF
           if done
              then return empty
              else do line <- getLine
                      let (c: dest: _) = split ": " line
                      let n = parseLine c $ split " " dest
                      unionWith (++) n <$> parse


getDistances' :: Int -> [Component] -> Network -> Map Component Int -> Map Component Int
getDistances' _ [] _ cd = cd
getDistances' d dest n cd = let currentDistances = union cd $ fromList [(c, d) | c <- dest]
                                next = nub (foldl (\ l c -> l ++ (n ! c)) [] dest) \\ keys currentDistances
                            in getDistances' (d + 1) next n currentDistances

getDistances :: Component -> Network -> Map Component Int
getDistances c n = getDistances' 0 [c] n empty

getPath :: Component -> Network -> Map Component Int -> [Component]
getPath c n d | d ! c == 0 = [c]
              | otherwise  = let newC = head $ keys $ Data.Map.filterWithKey
                                                      (\ src dest -> c `elem` dest && d ! src < d ! c) n
                             in getPath newC n d ++ [c]

removePath :: [Component] -> Network -> Network
removePath [_] n = n
removePath (h1: h2: t) n = let n1 = delete h2 (n ! h1)
                               n2 = delete h1 (n ! h2)
                           in removePath (h2: t) (insert h1 n1 $ insert h2 n2 n)

removeNPaths :: Int -> Component -> Component -> Network -> Network
removeNPaths 0 _ _ net = net
removeNPaths n c1 c2 net = let newNet = removePath (getPath c2 net (getDistances c1 net)) net
                           in removeNPaths (n - 1) c1 c2 newNet

iterateUntilPartition :: [Component] -> Network -> (Int, Int)
iterateUntilPartition (h1: h2: t) n = let removed = removeNPaths 3 h1 h2 n
                                          d1 = getDistances h1 removed
                                          d2 = getDistances h2 removed
                                      in if member h2 d1 then iterateUntilPartition (h1: t) n
                                         else (size d1, size d2)