module Lib
    ( runBFS
    ) where

import System.Environment ()
import Data.List as List ( map )
import Data.Map as Map ( Map, insert, findWithDefault, singleton, toList )
import Data.Set as Set ( Set, null, toList, empty, difference, unions, union, singleton )
import Control.Parallel.Strategies( using, rdeepseq, parBuffer )
import Control.DeepSeq ( force )

runBFS :: Map String (Set String) -> String -> String -> IO ()
runBFS neighborMap startWord endWord = do
  distances <- explore 0 (Set.singleton startWord) neighborMap (Map.singleton startWord 0) Set.empty
  let maxDegree = maximum $ List.map snd (Map.toList (force distances))
  -- Divide results by two since we are counting movies as vertices
  putStrLn ("Distance from " ++ startWord ++ " to " ++ endWord ++ ": " ++ show (toInteger (Map.findWithDefault (-1) endWord distances `div` 2) ))
  putStrLn ("Max distance: " ++ show (toInteger (maxDegree `div` 2)))

explore :: Int -> Set String -> Map String (Set String) -> Map String Int-> Set String -> IO (Map String Int)
explore level vertices neighborMap distances visited
    | Set.null vertices = return distances
    | otherwise = do
        let newVisited = Set.union visited vertices
        next <- getNeighborsForSet vertices neighborMap visited
        let newDistances = foldl (\d vertex -> insert vertex (level + 1) d) distances next
        explore (level + 1) next neighborMap newDistances newVisited

getNeighborsForSet :: Set String -> Map String (Set String) -> Set String -> IO (Set String)
getNeighborsForSet vertices neighborMap visited = do
    let vertexList = Set.toList vertices
    let neighborSets = List.map (getNeighbors neighborMap visited) vertexList `using` parBuffer 5000 rdeepseq
    -- let neighborSets = List.map (getNeighbors neighborMap visited) vertexList
    return (Set.unions neighborSets)

getNeighbors :: Map String (Set String) -> Set String -> String -> Set String
getNeighbors neighborMap visited vertex = Set.difference (Map.findWithDefault Set.empty vertex neighborMap) visited
