{-# language ViewPatterns #-}
module Agda.Utils.Graph.TopSort
    ( topSort
    ) where

import Data.List
import Data.Maybe
import Data.Function
import qualified Data.Set as Set
import qualified Data.Map as Map
import Control.Arrow

mergeBy :: (a -> a -> Bool) -> [a] -> [a] -> [a]
mergeBy _ [] xs = xs
mergeBy _ xs [] = xs
mergeBy f (x:xs) (y:ys)
    | f x y = x: mergeBy f xs (y:ys)
    | otherwise = y: mergeBy f (x:xs) ys

-- | topoligical sort with smallest-numbered available vertex first
-- | input: nodes, edges
-- | output is Nothing if the graph is not a DAG
topSort :: Ord n => [n] -> [(n, n)] -> Maybe [n]
topSort nodes edges = mergeBy (<) nodes' <$> g m is
  where
    g m []
        | Map.null m = Just []
        | otherwise = Nothing
    g m ((n, Set.toList -> cs): ns)
        = (n:) <$> g m' (mergeBy ((<) `on` fst) ns [(c, s) | c<-cs, (0, s) <- maybeToList $ Map.lookup c m'])
      where
        m' = foldr (Map.adjust $ first (+(-1))) (Map.delete n m) cs

    is = map (second snd) . filter ((==0) . fst . snd) $ Map.toList m

    nodes' = Set.toList $ Set.fromList nodes `Set.difference` Set.fromList (concatMap (\(a,b)->[a,b]) edges)

    m = foldr f mempty $ nub edges
    f (b, a)
        = Map.alter (Just . maybe (1, mempty) (first (+1))) b
        . Map.alter (Just . maybe (0, Set.singleton b) (second $ Set.insert b)) a