{-# LANGUAGE RecordWildCards #-}
{-# OPTIONS_HADDOCK hide #-}

-- | Implementation of Strongly Connected Components calculation. Use `AtCoder.Scc` instead.
--
-- @since 1.0.0.0
module AtCoder.Internal.Scc
  ( -- * Internal SCC
    SccGraph (nScc),

    -- * Constructors
    new,

    -- * Adding edges
    addEdge,

    -- * SCC calculation
    sccIds,
    scc,

    -- ** (Extra API) CSR API
    sccIdsCsr,
    sccCsr,
  )
where

import AtCoder.Internal.Csr qualified as ACICSR
import AtCoder.Internal.GrowVec qualified as ACIGV
import Control.Monad (unless, when)
import Control.Monad.Fix (fix)
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST, runST)
import Data.Foldable (for_)
import Data.Maybe (fromJust)
import Data.Vector qualified as V
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM

-- | Graph for collecting strongly connected components.
--
-- @since 1.0.0.0
data SccGraph s = SccGraph
  { -- | The number of vertices.
    --
    -- @since 1.0.0.0
    nScc :: {-# UNPACK #-} !Int,
    edgesScc :: !(ACIGV.GrowVec s (Int, Int))
  }

-- | \(O(n)\) Creates a `SccGraph` of \(n\) vertices.
--
-- @since 1.0.0.0
{-# INLINE new #-}
new :: (PrimMonad m) => Int -> m (SccGraph (PrimState m))
new nScc = do
  edgesScc <- ACIGV.new 0
  pure SccGraph {..}

-- | \(O(1)\) amortized. Adds an edge to the graph.
--
-- @since 1.0.0.0
{-# INLINE addEdge #-}
addEdge :: (PrimMonad m) => SccGraph (PrimState m) -> Int -> Int -> m ()
addEdge SccGraph {edgesScc} from to = do
  ACIGV.pushBack edgesScc (from, to)

-- | \(O(n + m)\) Returns a pair of @(# of scc, scc id)@.
--
-- @since 1.0.0.0
{-# INLINE sccIds #-}
sccIds :: (PrimMonad m) => SccGraph (PrimState m) -> m (Int, VU.Vector Int)
sccIds SccGraph {..} = do
  csr <- ACICSR.build' nScc <$> ACIGV.unsafeFreeze edgesScc
  pure $ sccIdsCsr csr

-- | \(O(n + m)\) Returns the strongly connected components.
--
-- @since 1.0.0.0
{-# INLINE scc #-}
scc :: (PrimMonad m) => SccGraph (PrimState m) -> m (V.Vector (VU.Vector Int))
scc g = stToPrim $ sccST g

-- -------------------------------------------------------------------------------------------------
-- Internal
-- -------------------------------------------------------------------------------------------------

{-# INLINEABLE sccST #-}
sccST :: SccGraph s -> ST s (V.Vector (VU.Vector Int))
sccST g = do
  (!groupNum, !ids) <- sccIds g
  let counts = VU.create $ do
        vec <- VUM.replicate groupNum (0 :: Int)
        VU.forM_ ids $ \x -> do
          VGM.modify vec (+ 1) x
        pure vec
  groups <- V.mapM VUM.unsafeNew $ VU.convert counts
  is <- VUM.replicate groupNum (0 :: Int)
  VU.iforM_ ids $ \v sccId -> do
    i <- VGM.read is sccId
    VGM.write is sccId $ i + 1
    VGM.write (groups VG.! sccId) i v
  V.mapM VU.unsafeFreeze groups

-- | \(O(n + m)\) API) Returns a pair of @(# of scc, scc id)@.
--
-- @since 1.1.0.0
{-# INLINEABLE sccIdsCsr #-}
sccIdsCsr :: ACICSR.Csr w -> (Int, VU.Vector Int)
sccIdsCsr g@ACICSR.Csr {..} = runST $ do
  -- see also the Wikipedia: https://en.wikipedia.org/wiki/Tarjan%27s_strongly_connected_components_algorithm#The_algorithm_in_pseudocode
  -- next SCC ID
  groupNum <- VUM.replicate 1 (0 :: Int)
  -- stack of vertices
  visited <- ACIGV.new nCsr
  -- vertex -> low-link: the smallest index of any node on the stack known to be reachable from
  -- v through v's DFS subtree, including v itself.
  low <- VUM.replicate nCsr (0 :: Int)
  -- vertex -> order of the visit (0, 1, ..)
  ord <- VUM.replicate nCsr (-1 :: Int)
  -- vertex -> scc id
  ids <- VUM.replicate nCsr (0 :: Int)

  let dfs v ord0 = do
        VGM.write low v ord0
        VGM.write ord v ord0
        ACIGV.pushBack visited v
        -- look around @v@, folding their low-link onto the low-link of @v@.
        ord' <-
          VU.foldM'
            ( \curOrd to -> do
                ordTo <- VGM.read ord to
                if ordTo == -1
                  then do
                    -- not visited yet.
                    nextOrd <- dfs to $ curOrd
                    lowTo <- VGM.read low to
                    VGM.modify low (min lowTo) v
                    pure nextOrd
                  else do
                    -- lookup back and update the low-link.
                    VGM.modify low (min ordTo) v
                    pure curOrd
            )
            (ord0 + 1)
            (g `ACICSR.adj` v)

        lowV <- VGM.read low v
        ordV <- VGM.read ord v
        when (lowV == ordV) $ do
          -- it's the root of a SCC, no more to look back
          sccId <- VGM.unsafeRead groupNum 0
          fix $ \loop -> do
            u <- fromJust <$> ACIGV.popBack visited
            VGM.write ord u nCsr
            VGM.write ids u sccId
            unless (u == v) loop
          VGM.unsafeWrite groupNum 0 $ sccId + 1
        pure ord'

  VU.foldM'_
    ( \curOrd i -> do
        o <- VGM.read ord i
        if o == -1
          then dfs i curOrd
          else pure curOrd
    )
    (0 :: Int)
    (VU.generate nCsr id)

  num <- VGM.unsafeRead groupNum 0
  -- The SCCs are reverse topologically sorted, e.g., [0, 1] <- [2] <- [3]
  -- Now reverse the SCC IDs so that they will be topologically sorted: [3] -> [2] -> [0, 1]
  for_ [0 .. nCsr - 1] $ \i -> do
    VGM.modify ids ((num - 1) -) i

  ids' <- VU.unsafeFreeze ids
  pure (num, ids')

-- | \(O(n + m)\) Returns the strongly connected components.
--
-- @since 1.1.0.0
{-# INLINEABLE sccCsr #-}
sccCsr :: ACICSR.Csr w -> V.Vector (VU.Vector Int)
sccCsr g = runST $ do
  groups <- V.mapM VUM.unsafeNew $ VU.convert counts
  is <- VUM.replicate groupNum (0 :: Int)
  VU.iforM_ ids $ \v sccId -> do
    i <- VGM.read is sccId
    VGM.write is sccId $ i + 1
    VGM.write (groups VG.! sccId) i v
  V.mapM VU.unsafeFreeze groups
  where
    (!groupNum, !ids) = sccIdsCsr g
    counts = VU.create $ do
      vec <- VUM.replicate groupNum (0 :: Int)
      VU.forM_ ids $ \x -> do
        VGM.modify vec (+ 1) x
      pure vec