import Control.Monad
import Control.Monad.Fix
import Control.Monad.Primitive
import Control.Monad.ST
import Data.Primitive.MutVar
import qualified Data.Vector as V
import qualified Data.Vector.Generic.Mutable as VM
import qualified Data.Vector.Unboxed.Mutable as VUM
import qualified Data.Vector.Unboxed as VU
import Data.List

newtype LVector s a = LVector (VU.MVector s a, MutVar s Int)

new :: (VU.Unbox a, PrimMonad m) => Int -> m (LVector (PrimState m) a)
new n = do
  vec <- VUM.new n
  len <- newMutVar 0
  return $ LVector (vec, len)

freeze
  :: (VU.Unbox a, PrimMonad m) => LVector (PrimState m) a -> m (VU.Vector a)
freeze (LVector (v, mv)) = do
  m <- readMutVar mv
  fmap (VU.take m) $ VU.freeze v

push :: (VU.Unbox a, PrimMonad m) => LVector (PrimState m) a -> a -> m ()
push lv@(LVector (vec, mv)) a = do
  m <- readMutVar mv
  when (VUM.length vec == m) $ do
    VUM.grow vec (VUM.length vec)
    return ()

  VUM.write vec m a
  writeMutVar mv (m + 1)

-------

type Graph node = [(node,[node])]

construct
  :: Eq node => Graph node -> (V.Vector node, GraphInternal, GraphInternal)
construct gnode = runST $ do
  let ns = V.fromList $ map fst gnode
  let rs = V.foldl' (\acc (v, k) i -> if i == k then v else acc i) (\_ -> (-1))
        $ V.indexed ns

  g  <- VM.replicateM (V.length ns) (new 10)
  gt <- VM.replicateM (V.length ns) (new 10)

  flip mapM_ (zip [0 ..] gnode) $ \(i, (_, es)) -> do
    flip mapM_ (map rs es) $ \e -> do
      gi <- VM.read g i
      push gi e

      gti <- VM.read gt e
      push gti i

  (,,)
    <$> pure ns
    <*> (V.mapM freeze =<< V.freeze g)
    <*> (V.mapM freeze =<< V.freeze gt)


type GraphInternal = V.Vector (VU.Vector Int)

scc :: Eq node => Graph node -> V.Vector (V.Vector node)
scc gr = V.map (V.map (mapper V.!) . V.convert) build
 where
  (mapper, g, gt) = construct gr

  order :: VU.Vector Int
  order = VU.reverse $ runST $ do
    used  <- VUM.replicate (V.length g) False
    order <- new 10

    flip fix (0 :: Int) $ \f i -> do
      b <- VM.read used i
      if b
        then return (-1)
        else do
          VM.write used i True
          VU.mapM_ f     (g V.! i)

          push     order i

          if i == V.length g - 1 then return 0 else f (i + 1)

    freeze order

  computed :: VU.Vector Int
  size :: Int
  (computed, size) = runST $ do
    comp <- VUM.replicate (V.length g) (-1)
    cnt  <- newMutVar 0

    VU.forM_ order $ \o -> do
      b <- VUM.read comp o
      c <- readMutVar cnt

      when (b == -1) $ do
        go comp o c
        writeMutVar cnt (c + 1)

    (,) <$> VU.freeze comp <*> readMutVar cnt
   where
    go comp idx cnt = do
      b <- VUM.read comp idx
      when (b == -1) $ do
        VUM.write comp idx cnt
        VU.mapM_ (\to -> go comp to cnt) (gt V.! idx)

  build :: GraphInternal
  build = runST $ do
    comps <- VM.replicateM size $ new 10

    forM_ [0 .. V.length g - 1] $ \i -> do
      VU.forM_ (g V.! i) $ \to -> do
        let x = computed VU.! i
        let y = computed VU.! to
        when (x == y) $ do
          tx <- VM.read comps x
          push tx i

    V.mapM freeze =<< V.freeze comps

example :: Graph String
example =
  [ ("1", ["2"])
  , ("2", ["3", "5", "6"])
  , ("3", ["4", "7"])
  , ("4", ["3", "8"])
  , ("5", ["1", "6"])
  , ("6", ["7"])
  , ("7", ["6", "8"])
  , ("8", ["8"])
  ]

main = print $ scc example