content: -- | Internal implementation of @AtCoder.String@ module.
  module AtCoder.Internal.String
    ( -- * Suffix array
      saNaive,
      saDoubling,
      saIsImpl,
      saIs,
      saIsManual,
    )
  where
  
  import Control.Monad (unless, when)
  import Control.Monad.ST (runST)
  import Data.Bit (Bit (..))
  import Data.Foldable (for_)
  import Data.Vector.Algorithms.Intro qualified as VAI
  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
  import GHC.Stack (HasCallStack)
  
  -- TODO: remove `HasCallStack`?
  
  -- | \(O(n^2)\) Internal implementation of suffix array creation (naive).
  --
  -- @since 1.0.0.0
  {-# INLINEABLE saNaive #-}
  saNaive :: (HasCallStack) => VU.Vector Int -> VU.Vector Int
  saNaive s =
    let n = VU.length s
        f l0 r0
          | l0 == r0 = GT
          | otherwise = inner l0 r0
          where
            inner l r
              | l < n && r < n =
                  let sl = s VG.! l
                      sr = s VG.! r
                   in if sl /= sr
                        then compare sl sr
                        else inner (l + 1) (r + 1)
              | l == n = LT
              | otherwise = GT
     in -- modify + generate should fuse
        VU.modify (VAI.sortBy f) $ VU.generate n id
  
  -- | \(O(n \log n)\) Internal implementation of suffix array creation (doubling).
  --
  -- @since 1.0.0.0
  {-# INLINEABLE saDoubling #-}
  saDoubling :: (HasCallStack) => VU.Vector Int -> VU.Vector Int
  saDoubling s = VU.create $ do
    let n = VU.length s
    sa <- VUM.generate n id
    let loop rnk tmp k
          | k >= n = pure ()
          | otherwise = do
              rnk' <- VU.unsafeFreeze rnk
              let cmp x y =
                    let rnkX = rnk' VG.! x
                        rnkY = rnk' VG.! y
                     in if rnkX /= rnkY
                          then compare rnkX rnkY
                          else
                            let rx = if x + k < n then rnk' VG.! (x + k) else (-1)
                                ry = if y + k < n then rnk' VG.! (y + k) else (-1)
                             in compare rx ry
              VAI.sortBy cmp sa
              sa' <- VU.unsafeFreeze sa
              VGM.write tmp (sa' VG.! 0) 0
              VU.zipWithM_
                ( \saI_ saI -> do
                    tmpI_ <- VGM.read tmp saI_
                    VGM.write tmp saI $! tmpI_ + if cmp saI_ saI == LT then 1 else 0
                )
                sa'
                (VU.tail sa')
              loop tmp rnk (2 * k)
    rnkVec <- VU.thaw s
    tmpVec <- VUM.unsafeNew n
    loop rnkVec tmpVec 1
    pure sa
  
  -- TODO: use backpermute
  
  -- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting).
  --
  -- @since 1.0.0.0
  {-# INLINEABLE saIsImpl #-}
  saIsImpl ::
    (HasCallStack) =>
    -- | naive threshold
    Int ->
    -- | doubling threshold
    Int ->
    -- | string
    VU.Vector Int ->
    -- | upper bounds
    Int ->
    -- | suffix array
    VU.Vector Int
  saIsImpl naiveThreshold doublingThreshold s upper = VU.create $ do
    let n = VU.length s
    let !ls = VU.create $ do
          ls_ <- VUM.unsafeNew @_ @Bit n
          VGM.write ls_ (n - 1) $ Bit False
          for_ [n - 2, n - 3 .. 0] $ \i -> do
            let si = s VG.! i
            let si1 = s VG.! (i + 1)
            b <-
              if si == si1
                then VGM.read ls_ (i + 1)
                else pure . Bit $ si < si1
            VGM.write ls_ i b
          pure ls_
  
    let (!sumL, !sumS) = runST $ do
          sumL_ <- VUM.replicate (upper + 1) (0 :: Int)
          sumS_ <- VUM.replicate (upper + 1) (0 :: Int)
          VU.zipWithM_
            ( \(Bit !b) !si -> do
                if not b
                  then VGM.modify sumS_ (+ 1) si
                  else VGM.modify sumL_ (+ 1) $ si + 1
            )
            ls
            s
          for_ [0 .. upper] $ \i -> do
            l <- VGM.read sumL_ i
            VGM.modify sumS_ (+ l) i
            when (i < upper) $ do
              si <- VGM.read sumS_ i
              VGM.modify sumL_ (+ si) $ i + 1
          sumL' <- VU.unsafeFreeze sumL_
          sumS' <- VU.unsafeFreeze sumS_
          pure (sumL', sumS')
  
    sa <- VUM.unsafeNew @_ @Int n
    let induce lms = do
          VGM.set sa (-1)
          buf <- VU.thaw sumS
          VU.forM_ lms $ \d -> do
            unless (d == n) $ do
              i <- VGM.read buf $ s VG.! d
              VGM.write buf (s VG.! d) $ i + 1
              VGM.write sa i d
          VU.iforM_ sumL (VGM.write buf)
          do
            i <- VGM.read buf $ s VG.! (n - 1)
            VGM.write buf (s VG.! (n - 1)) $ i + 1
            VGM.write sa i $ n - 1
          do
            -- TODO: try VUM.forM_
            for_ [0 .. n - 1] $ \i -> do
              v <- VGM.read sa i
              when (v >= 1 && not (unBit (ls VG.! (v - 1)))) $ do
                j <- VGM.read buf $ s VG.! (v - 1)
                VGM.write buf (s VG.! (v - 1)) $ j + 1
                VGM.write sa j $ v - 1
          VU.iforM_ sumL (VGM.write buf)
          -- TODO: try foldr
          for_ [n - 1, n - 2 .. 0] $ \i -> do
            v <- VGM.read sa i
            when (v >= 1 && unBit (ls VG.! (v - 1))) $ do
              j <- subtract 1 <$> VGM.read buf (s VG.! (v - 1) + 1)
              VGM.write buf (s VG.! (v - 1) + 1) j
              VGM.write sa j $ v - 1
  
    let (!lmsMap, !m) = runST $ do
          lmsMap_ <- VUM.replicate (n + 1) (-1 :: Int)
          len <-
            VU.ifoldM'
              ( \iAcc i (Bit !b1, Bit !b2) -> do
                  if not b1 && b2
                    then do
                      VGM.write lmsMap_ (i + 1) iAcc
                      pure $ iAcc + 1
                    else do
                      pure iAcc
              )
              (0 :: Int)
              $ VU.zip ls (VU.tail ls)
          (,len) <$> VU.unsafeFreeze lmsMap_
  
    let lms = VU.create $ do
          lms_ <- VUM.unsafeNew @_ @Int m
          VU.ifoldM'_
            ( \iAcc i (Bit !b1, Bit !b2) -> do
                if not b1 && b2
                  then do
                    VGM.write lms_ iAcc $ i + 1
                    pure $ iAcc + 1
                  else do
                    pure iAcc
            )
            (0 :: Int)
            $ VU.zip ls (VU.tail ls)
          pure lms_
  
    induce lms
  
    when (m > 0) $ do
      sa' <- VU.unsafeFreeze sa
      sortedLms <- VUM.unsafeNew @_ @Int m
      VU.foldM'_
        ( \iAcc v -> do
            if lmsMap VG.! v /= -1
              then do
                VGM.write sortedLms iAcc v
                pure $ iAcc + 1
              else do
                pure iAcc
        )
        (0 :: Int)
        sa'
      recS <- VUM.unsafeNew @_ @Int m
      (\i -> VGM.write recS i 0) . (lmsMap VG.!) =<< VGM.read sortedLms 0
  
      recUpper' <- do
        sortedLms' <- VU.unsafeFreeze sortedLms
        VU.foldM'
          ( \recUpper (!l, !r) -> do
              let lmsMapL = lmsMap VG.! l
              let lmsMapR = lmsMap VG.! r
              let endL = if lmsMapL + 1 < m then lms VG.! (lmsMapL + 1) else n
              let endR = if lmsMapR + 1 < m then lms VG.! (lmsMapR + 1) else n
              let same
                    | endL - l /= endR - r = False
                    | otherwise = inner l r
                    where
                      inner x y
                        | x < endL && s VG.! x == s VG.! y = inner (x + 1) (y + 1)
                        | x == n || s VG.! x /= s VG.! y = False
                        | otherwise = True
              let recUpper' = if same then recUpper else recUpper + 1
              VGM.write recS lmsMapR recUpper'
              pure recUpper'
          )
          (0 :: Int)
          $ VU.zip sortedLms' (VU.tail sortedLms')
  
      recS' <- VU.unsafeFreeze recS
      let recSa = saIsManual naiveThreshold doublingThreshold recS' recUpper'
  
      VU.iforM_ recSa $ \i x -> do
        VGM.write sortedLms i (lms VG.! x)
      induce =<< VU.unsafeFreeze sortedLms
  
    pure sa
  
  -- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting).
  --
  -- SA-IS, linear-time suffix array construction.
  -- Reference:
  -- G. Nong, S. Zhang, and W. H. Chan,
  -- Two Efficient Algorithms for Linear Time Suffix Array Construction
  --
  -- @since 1.0.0.0
  {-# INLINE saIs #-}
  saIs ::
    (HasCallStack) =>
    -- | string
    VU.Vector Int ->
    -- | upper bounds
    Int ->
    -- | suffix array
    VU.Vector Int
  saIs = saIsManual 10 40
  
  -- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting).
  --
  -- SA-IS, linear-time suffix array construction.
  -- Reference:
  -- G. Nong, S. Zhang, and W. H. Chan,
  -- Two Efficient Algorithms for Linear Time Suffix Array Construction
  --
  -- @since 1.0.0.0
  {-# INLINE saIsManual #-}
  saIsManual ::
    (HasCallStack) =>
    -- | naive threshold
    Int ->
    -- | doubling threshold
    Int ->
    -- | string
    VU.Vector Int ->
    -- | upper bounds
    Int ->
    -- | suffix array
    VU.Vector Int
  saIsManual naiveThreshold doublingThreshold s upper
    | n == 0 = VU.empty
    | n == 1 = VU.singleton 0
    | n == 2 && s VG.! 0 < s VG.! 1 = VU.fromListN 2 [0, 1]
    | n == 2 = VU.fromListN 2 [1, 0]
    | n < naiveThreshold = saNaive s
    | n < doublingThreshold = saDoubling s
    | otherwise = saIsImpl naiveThreshold doublingThreshold s upper
    where
      n = VU.length s
  