content: {-# LANGUAGE RecordWildCards #-}
  
  -- | It solves [Minimum-cost flow problem](https://en.wikipedia.org/wiki/Minimum-cost_flow_problem).
  --
  -- ==== __Example__
  -- Create min cost flow graph (`McfGraph`):
  --
  -- >>> import AtCoder.MinCostFlow qualified as MCF
  -- >>> g <- MCF.new @_ {- capacity -} @Int {- cost -} @Int 4
  --
  -- Build a simple graph with @'addEdge' g from to cap cost@ or `addEdge_`:
  --
  -- >>> MCF.addEdge g 0 1 2 3    --  0 --> 1     2
  -- 0
  --
  -- >>> MCF.addEdge_ g 1 2 2 5   --  0 --> 1 --> 2
  --
  -- Augment flow with `flow`, `maxFlow` or `slope`:
  --
  -- >>> MCF.slope g 0 2 maxBound -- slope g from to flowLimit
  -- [(0,0),(2,16)]
  --
  -- Note that you can't call `flow`, `maxFlow` or `slope` multiple times, or else you'll get wrong
  -- return value.
  --
  -- @since 1.0.0.0
  module AtCoder.MinCostFlow
    ( -- * Minimum cost flow
      McfGraph (nG),
  
      -- * Constructor
      new,
  
      -- * Graph building
      addEdge,
      addEdge_,
  
      -- * Flow and slope
      flow,
      maxFlow,
      slope,
  
      -- * Edge information
      getEdge,
      edges,
      unsafeEdges,
    )
  where
  
  -- TODO: add `maxCostFlow`.
  -- TODO: add `build`.
  -- TODO: is this fast enough with `INLINEABLE`?
  
  import AtCoder.Internal.Assert qualified as ACIA
  import AtCoder.Internal.Buffer qualified as ACIB
  import AtCoder.Internal.GrowVec qualified as ACIGV
  import AtCoder.Internal.McfCsr qualified as ACIMCSR
  import AtCoder.Internal.MinHeap qualified as ACIMH
  import Control.Monad (unless, when)
  import Control.Monad.Fix (fix)
  import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
  import Control.Monad.ST (ST)
  import Data.Bit (Bit (..))
  import Data.Maybe (fromJust)
  import Data.Primitive.MutVar (readMutVar)
  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)
  
  -- | Min cost flow graph.
  --
  -- @since 1.0.0.0
  data McfGraph s cap cost = McfGraph
    { -- | The number of vertices.
      --
      -- @since 1.0.0.0
      nG :: {-# UNPACK #-} !Int,
      -- | fromVertex -> vector of @(from, to, cap, flow, cost)@.
      edgesG :: !(ACIGV.GrowVec s (Int, Int, cap, cap, cost))
    }
  
  -- | Creates a directed graph with \(n\) vertices and \(0\) edges. @cap@ and @cost@ are the type of
  -- the capacity and the cost, respectively.
  --
  -- ==== Constraints
  -- - \(0 \leq n\)
  --
  -- ==== Complexity
  -- - \(O(n)\)
  --
  -- @since 1.0.0.0
  {-# INLINE new #-}
  new :: (PrimMonad m, VU.Unbox cap, VU.Unbox cost) => Int -> m (McfGraph (PrimState m) cap cost)
  new nG = stToPrim $ do
    edgesG <- ACIGV.new 0
    pure McfGraph {..}
  
  -- | Adds an edge oriented from @from@ to @to@ with capacity @cap@ and cost @cost@. It returns an
  -- integer \(k\) such that this is the \(k\)-th edge that is added.
  --
  -- ==== Constraints
  -- - \(0 \leq \mathrm{from}, \mathrm{to} \lt n\)
  -- - \(0 \leq \mathrm{cap}, \mathrm{cost}\)
  --
  -- ==== Complexity
  -- - \(O(1)\) amortized
  --
  -- @since 1.0.0.0
  {-# INLINE addEdge #-}
  addEdge ::
    (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph (PrimState m) cap cost ->
    -- | from
    Int ->
    -- | to
    Int ->
    -- | capacity
    cap ->
    -- | cost
    cost ->
    -- | Edge index
    m Int
  addEdge g from to cap cost = stToPrim $ addEdgeST g from to cap cost
  
  -- | `addEdge` with the return value discarded.
  --
  -- ==== Constraints
  -- - \(0 \leq \mathrm{from}, \mathrm{to} \lt n\)
  -- - \(0 \leq \mathrm{cap}, \mathrm{cost}\)
  --
  -- ==== Complexity
  -- - \(O(1)\) amortized
  --
  -- @since 1.0.0.0
  {-# INLINE addEdge_ #-}
  addEdge_ ::
    (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph (PrimState m) cap cost ->
    -- | from
    Int ->
    -- | to
    Int ->
    -- | capacity
    cap ->
    -- | cost
    cost ->
    m ()
  addEdge_ graph from to cap cost = stToPrim $ do
    _ <- addEdgeST graph from to cap cost
    pure ()
  
  -- | Augments the flow from \(s\) to \(t\) as much as possible, until reaching the amount of
  -- @flowLimit@. It returns the amount of the flow and the cost.
  --
  -- ==== Constraints
  -- - Same as `slope`.
  --
  -- ==== Complexity
  -- - Same as `slope`.
  --
  -- @since 1.0.0.0
  {-# INLINE flow #-}
  flow ::
    (HasCallStack, PrimMonad m, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph (PrimState m) cap cost ->
    -- | Source @s@
    Int ->
    -- | Sink @t@
    Int ->
    -- | Flow limit
    cap ->
    -- | Tuple of @(cap, cost@)
    m (cap, cost)
  flow graph s t flowLimit = stToPrim $ do
    res <- slopeST graph s t flowLimit
    pure $ VG.last res
  
  -- | `flow` with no capacity limit.
  --
  -- ==== Constraints
  -- - Same as `slope`.
  --
  -- ==== Complexity
  -- - Same as `slope`.
  --
  -- @since 1.0.0.0
  {-# INLINE maxFlow #-}
  maxFlow ::
    (HasCallStack, PrimMonad m, Integral cap, Ord cap, Bounded cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph (PrimState m) cap cost ->
    -- | Source @s@
    Int ->
    -- | Sink @t@
    Int ->
    -- | Tuple of @(cap, cost@)
    m (cap, cost)
  maxFlow graph s t = stToPrim $ do
    res <- slopeST graph s t maxBound
    pure $ VG.last res
  
  -- | Let \(g\) be a function such that \(g(x)\) is the cost of the minimum cost \(s-t\) flow when
  -- the amount of the flow is exactly \(x\). \(g\) is known to be piecewise linear.
  --
  -- - The first element of the list is \((0, 0)\).
  -- - Both of first and second tuple elements are strictly increasing.
  -- - No three changepoints are on the same line.
  -- - The last element of the list is \(y, g(y))\), where \(y = \min(x, \mathrm{flowLimit})\).
  --
  -- ==== Constraints
  --
  -- Let \(x\) be the maximum cost among all edges.
  --
  -- - \(s \neq t\)
  -- - \(0 \leq s, t \lt n\)
  -- - You can't call `slope`, `flow` or `maxFlow` multiple times.
  -- - The total amount of the flow is in @cap@.
  -- - The total cost of the flow is in @cost@.
  -- - \(0 \leq nx \leq 8 \times 10^{18} + 1000\)
  --
  -- ==== Complexity
  -- - \(O(F (n + m) \log (n + m))\), where \(F\) is the amount of the flow and \(m\) is the number of added edges.
  --
  -- @since 1.0.0.0
  {-# INLINE slope #-}
  slope ::
    (HasCallStack, PrimMonad m, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph (PrimState m) cap cost ->
    -- | Source @s@
    Int ->
    -- | Sink @t@
    Int ->
    -- | Flow limit
    cap ->
    -- | Vector of @(cap, cost)@
    m (VU.Vector (cap, cost))
  slope g s t flowLimit = stToPrim $ slopeST g s t flowLimit
  
  -- | Returns the current internal state of the edges: @(from, to, cap, flow, cost)@. The edges are
  -- ordered in the same order as added by `addEdge`.
  --
  -- ==== Constraints
  -- - \(0 \leq i \lt m\)
  --
  -- ==== Complexity
  -- - \(O(1)\)
  --
  -- @since 1.0.0.0
  {-# INLINE getEdge #-}
  getEdge ::
    (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph (PrimState m) cap cost ->
    -- | Edge index
    Int ->
    -- | Tuple of @(from, to, cap, flow, cost)@
    m (Int, Int, cap, cap, cost)
  getEdge g i = stToPrim $ getEdgeST g i
  
  -- | Returns the current internal state of the edges: @(from, to, cap, flow, cost)@. The edges are
  -- ordered in the same order as added by `addEdge`.
  --
  -- ==== Complexity
  -- - \(O(m)\), where \(m\) is the number of added edges.
  --
  -- @since 1.0.0.0
  {-# INLINE edges #-}
  edges ::
    (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph (PrimState m) cap cost ->
    -- | Vector of @(from, to, cap, flow, cost)@
    m (VU.Vector (Int, Int, cap, cap, cost))
  edges McfGraph {..} = stToPrim $ do
    ACIGV.freeze edgesG
  
  -- | Returns the current internal state of the edges: @(from, to, cap, flow, cost)@, but without
  -- making copy. The edges are ordered in the same order as added by `addEdge`.
  --
  -- ==== Complexity
  -- - \(O(1)\)
  --
  -- @since 1.0.0.0
  {-# INLINE unsafeEdges #-}
  unsafeEdges ::
    (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph (PrimState m) cap cost ->
    -- | Vector of @(from, to, cap, flow, cost)@
    m (VU.Vector (Int, Int, cap, cap, cost))
  unsafeEdges McfGraph {..} = stToPrim $ do
    ACIGV.unsafeFreeze edgesG
  
  -- -------------------------------------------------------------------------------------------------
  -- Internal
  -- -------------------------------------------------------------------------------------------------
  
  {-# INLINEABLE addEdgeST #-}
  addEdgeST ::
    (HasCallStack, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph s cap cost ->
    -- | from
    Int ->
    -- | to
    Int ->
    -- | capacity
    cap ->
    -- | cost
    cost ->
    -- | Edge index
    ST s Int
  addEdgeST McfGraph {..} from to cap cost = do
    let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.addEdgeST" "`from` vertex" from "the number of vertices" nG
    let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.addEdgeST" "`to` vertex" to "the number of vertices" nG
    let !_ = ACIA.runtimeAssert (0 <= cap) "AtCoder.MinCostFlow.addEdgeST: given invalid edge `cap` less than `0`"
    let !_ = ACIA.runtimeAssert (0 <= cost) "AtCoder.MinCostFlow.addEdgeST: given invalid edge `cost` less than `0`"
    m <- ACIGV.length edgesG
    ACIGV.pushBack edgesG (from, to, cap, 0, cost)
    pure m
  
  {-# INLINEABLE slopeST #-}
  slopeST ::
    (HasCallStack, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph s cap cost ->
    -- | Source @s@
    Int ->
    -- | Sink @t@
    Int ->
    -- | Flow limit
    cap ->
    -- | Vector of @(cap, cost)@
    ST s (VU.Vector (cap, cost))
  slopeST McfGraph {..} s t flowLimit = do
    let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.slopeST" "`source` vertex" s "the number of vertices" nG
    let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.slopeST" "`sink` vertex" t "the number of vertices" nG
    let !_ = ACIA.runtimeAssert (s /= t) "AtCoder.MinCostFlow.slopeST: `source` and `sink` vertex must be distinct"
  
    edges_@(VU.V_5 _ _ _ caps _ _) <- ACIGV.unsafeFreeze edgesG
    (!edgeIdx, !g) <- ACIMCSR.build nG edges_
    result <- internalSlopeST g nG s t flowLimit
  
    (VUM.MV_5 _ _ _ _ flows _) <- readMutVar $ ACIGV.vecGV edgesG
    VU.iforM_ (VU.zip caps edgeIdx) $ \v (!cap1, !iEdge) -> do
      cap2 <- VGM.read (ACIMCSR.capCsr g) iEdge
      VGM.write flows v $! cap1 - cap2
  
    pure result
  
  {-# INLINEABLE internalSlopeST #-}
  internalSlopeST ::
    forall cap cost s.
    (HasCallStack, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>
    ACIMCSR.Csr s cap cost ->
    Int ->
    Int ->
    Int ->
    cap ->
    ST s (VU.Vector (cap, cost))
  internalSlopeST csr@ACIMCSR.Csr {..} n s t flowLimit = do
    duals <- VUM.replicate n 0
    dists <- VUM.unsafeNew n :: ST s (VUM.MVector s cost)
    prevE <- VUM.unsafeNew n :: ST s (VUM.MVector s Int)
    vis <- VUM.unsafeNew n :: ST s (VUM.MVector s Bit)
  
    -- FIXME: maximum capacity of heap?
    let nEdges = VU.length toCsr
    queMin <- ACIB.new nEdges :: ST s (ACIB.Buffer s Int)
    heap <- ACIMH.new nEdges :: ST s (ACIMH.Heap s (cost, Int))
  
    let dualRef = do
          VGM.set dists $ maxBound @cost
          VGM.set vis $ Bit False
          ACIB.clear queMin
          -- TODO: compare with the first element only, so make up custom Q data type
          ACIMH.clear heap
  
          VGM.write dists s 0
          ACIB.pushBack queMin s
          fix $ \loop -> do
            b1 <- ACIB.null queMin
            b2 <- ACIMH.null heap
            when (not b1 || not b2) $ do
              v <-
                if not b1
                  then do
                    fromJust <$> ACIB.popBack queMin
                  else do
                    (!_, !to) <- fromJust <$> ACIMH.pop heap
                    pure to
  
              Bit visV <- VGM.exchange vis v $ Bit True
              unless (v == t) $ do
                unless visV $ do
                  -- dist[v] = shortest(s, v) + dual[s] - dual[v]
                  -- dist[v] >= 0 (all reduced cost are positive)
                  -- dist[v] <= (n-1)C
                  dualV <- VGM.read duals v
                  distV <- VGM.read dists v
                  let start = startCsr VG.! v
                  VU.iforM_ (ACIMCSR.adj csr v) $ \di (!to, !rev, !cost) -> do
                    cap <- VGM.read capCsr $ start + di
  
                    unless (cap == 0) $ do
                      -- - |-dual[e.to] + dual[v]| <= (n-1)C
                      -- - cost <= C - -(n-1)C + 0 = nC
                      cost' <- do
                        dualTo <- VGM.read duals to
                        pure $! cost - dualTo + dualV
                      distTo <- VGM.read dists to
                      when (distTo - distV > cost') $ do
                        let !distTo' = distV + cost'
                        VGM.write dists to distTo'
                        VGM.write prevE to rev
                        if distTo' == distV
                          then ACIB.pushBack queMin to
                          else ACIMH.push heap (distTo', to)
  
                loop
          Bit visT <- VGM.read vis t
          when visT $ do
            distT <- VGM.read dists t
            vis' <- VU.unsafeFreeze vis
            dists' <- VU.unsafeFreeze dists
            VU.iforM_ (VU.zip vis' dists') $ \v (Bit !visV, !distV) -> do
              when visV $ do
                VGM.modify duals (subtract (distT - distV)) v
          pure visT
  
    result <- ACIGV.new 16
    ACIGV.pushBack result (0 :: cap, 0 :: cost)
  
    let inner :: cap -> cost -> cost -> ST s ()
        inner flow_ cost prevCostPerFlow =
          when (flow_ < flowLimit) $ do
            b <- dualRef
            when b $ do
              prevE' <- VU.unsafeFreeze prevE
  
              let minC :: cap -> Int -> ST s cap
                  minC !acc v
                    | v == s = pure acc
                    | otherwise = do
                        let iPrev = prevE' VG.! v
                        cap <- VGM.read capCsr $ revCsr VG.! iPrev
                        minC (min acc cap) $ toCsr VG.! iPrev
              c <- minC (flowLimit - flow_) t
  
              let subC :: Int -> ST s ()
                  subC v = when (v /= s) $ do
                    let iPrev = prevE' VG.! v
                    VGM.modify capCsr (+ c) iPrev
                    VGM.modify capCsr (subtract c) $ revCsr VG.! iPrev
                    subC $ toCsr VG.! iPrev
              subC t
  
              d <- negate <$> VGM.read duals s
              let !flow' = flow_ + c
              let !cost' = cost + fromIntegral c * d -- TODO: minimize the type boundary
              when (prevCostPerFlow == d) $ do
                ACIGV.popBack_ result
              ACIGV.pushBack result (flow', cost')
              inner flow' cost' d
  
    inner 0 0 (-1)
    ACIGV.unsafeFreeze result
  
  {-# INLINEABLE getEdgeST #-}
  getEdgeST ::
    (HasCallStack, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
    -- | Graph
    McfGraph s cap cost ->
    -- | Edge index
    Int ->
    -- | Tuple of @(from, to, cap, flow, cost)@
    ST s (Int, Int, cap, cap, cost)
  getEdgeST McfGraph {..} i = do
    m <- ACIGV.length edgesG
    let !_ = ACIA.checkEdge "AtCoder.MinCostFlow.getEdgeST" i m
    ACIGV.read edgesG i
  