content: {-# LANGUAGE RecordWildCards #-}
  {-# OPTIONS_HADDOCK hide #-}
  
  -- | Internal CSR for `AtCoder.MinCostFlow`.
  --
  -- @since 1.0.0.0
  module AtCoder.Internal.McfCsr
    ( -- * Compressed sparse row
      Csr (..),
  
      -- * Constructor
      build,
  
      -- * Accessor
      adj,
    )
  where
  
  import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
  import Control.Monad.ST (ST)
  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.Base qualified as VU
  import Data.Vector.Unboxed.Mutable qualified as VUM
  import GHC.Stack (HasCallStack)
  
  -- | CSR for min cost flow.
  --
  -- @since 1.0.0.0
  data Csr s cap cost = Csr
    { -- | @since 1.0.0.0
      startCsr :: !(VU.Vector Int),
      -- | @since 1.0.0.0
      toCsr :: !(VU.Vector Int),
      -- | @since 1.0.0.0
      revCsr :: !(VU.Vector Int),
      -- | Mutable.
      --
      -- @since 1.0.0.0
      capCsr :: !(VUM.MVector s cap),
      -- | @since 1.0.0.0
      costCsr :: !(VU.Vector cost)
    }
  
  -- | \(O(n + m)\) Creates `Csr`.
  --
  -- @since 1.0.0.0
  {-# INLINE build #-}
  build :: (HasCallStack, PrimMonad m, Num cap, VU.Unbox cap, VU.Unbox cost, Num cost) => Int -> VU.Vector (Int, Int, cap, cap, cost) -> m (VU.Vector Int, Csr (PrimState m) cap cost)
  build n edges = stToPrim $ buildST n edges
  
  -- | \(O(1)\) Returns a vector of @(to, rev, cost)@.
  --
  -- @since 1.0.0.0
  {-# INLINE adj #-}
  adj :: (HasCallStack, Num cap, VU.Unbox cap, VU.Unbox cost) => Csr s cap cost -> Int -> VU.Vector (Int, Int, cost)
  adj Csr {..} v = VU.slice offset len vec
    where
      offset = startCsr VG.! v
      len = startCsr VG.! (v + 1) - offset
      vec = VU.zip3 toCsr revCsr costCsr
  
  {-# INLINEABLE buildST #-}
  buildST :: (HasCallStack, Num cap, VU.Unbox cap, VU.Unbox cost, Num cost) => Int -> VU.Vector (Int, Int, cap, cap, cost) -> ST s (VU.Vector Int, Csr s cap cost)
  buildST n edges = do
    let m = VU.length edges
    -- create the offsets first (this is a different step from ac-library)
    let startCsr = VU.create $ do
          start <- VUM.replicate (n + 1) (0 :: Int)
          -- count degrees
          let (VU.V_5 _ froms tos _ _ _) = edges
          VU.forM_ (VU.zip froms tos) $ \(!from, !to) -> do
            VGM.modify start (+ 1) $ from + 1
            VGM.modify start (+ 1) $ to + 1
          -- sum up the degrees
          VUM.iforM_ (VUM.init start) $ \i dx -> do
            VGM.modify start (+ dx) (i + 1)
          pure start
  
    toVec <- VUM.unsafeNew $ 2 * m
    revVec <- VUM.unsafeNew $ 2 * m
    capCsr <- VUM.unsafeNew $ 2 * m
    costVec <- VUM.unsafeNew $ 2 * m
  
    -- build CSR
    counter <- VU.thaw startCsr
    edgeIdx <- VU.forM edges $ \(!from, !to, !cap, !flow, !cost) -> do
      i1 <- VGM.read counter from
      VGM.modify counter (+ 1) from
      i2 <- VGM.read counter to
      VGM.modify counter (+ 1) to
      -- write forward edge
      VGM.write toVec i1 to
      VGM.write revVec i1 i2
      VGM.write capCsr i1 $! cap - flow
      VGM.write costVec i1 cost
      -- write backward edge
      VGM.write toVec i2 from
      VGM.write revVec i2 i1
      VGM.write capCsr i2 flow
      VGM.write costVec i2 (-cost)
      -- remember forward edge index
      pure i1
  
    toCsr <- VU.unsafeFreeze toVec
    revCsr <- VU.unsafeFreeze revVec
    costCsr <- VU.unsafeFreeze costVec
    pure (edgeIdx, Csr {..})
  