content: {-# LANGUAGE RecordWildCards #-}
  
  -- | Two-dimensional segment tree for commutative monoids at fixed points.
  --
  -- ==== `SegTree2d` vs `WaveletMatrix2d`
  -- They basically have the same functionalities and performance, however, `SegTree2d` performs better in
  -- @ac-library-hs@.
  --
  -- ==== __Examples__
  -- Create a two-dimensional segment tree for points \((0, 0)\) with weight \(10\) and \((1, 1)\)
  -- with weight \(20\):
  --
  -- >>> import AtCoder.Extra.SegTree2d qualified as Seg
  -- >>> import Data.Semigroup (Sum (..))
  -- >>> import Data.Vector.Unboxed qualified as VU
  -- >>> seg <- Seg.build3 @_ @(Sum Int) $ VU.fromList [(0, 0, 10), (1, 1, 20)]
  --
  -- Get monoid product in \([x_1, x_2) \times [y_1, y_2)\) with `prod`:
  --
  -- >>> Seg.prod seg {- x -} 0 2 {- y -} 0 2
  -- Sum {getSum = 30}
  --
  -- Monoid values can be altered:
  --
  -- >>> Seg.write seg 1 50
  -- >>> Seg.prod seg {- x -} 1 2 {- y -} 1 2
  -- Sum {getSum = 50}
  --
  -- >>> Seg.allProd seg
  -- Sum {getSum = 60}
  --
  -- >>> Seg.count seg {- x -} 0 2 {- y -} 0 2
  -- 2
  --
  -- @since 1.2.3.0
  module AtCoder.Extra.SegTree2d
    ( -- * SegTree2d
      SegTree2d (..),
  
      -- * Constructors
      new,
      build,
      build2,
      build3,
  
      -- * Write
  
      -- read,
      write,
      modify,
      modifyM,
  
      -- * Monoid product
      prod,
      allProd,
  
      -- * Count
      count,
    )
  where
  
  import AtCoder.Extra.Bisect (lowerBound)
  import AtCoder.Extra.Vector (argsort)
  import AtCoder.Internal.Assert qualified as ACIA
  import AtCoder.Internal.Bit qualified as ACIB
  import Control.Monad (when)
  import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
  import Control.Monad.ST (ST)
  import Data.Bits
  import Data.Foldable (for_)
  import Data.Maybe (fromMaybe)
  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)
  import Prelude hiding (read)
  
  -- | Two-dimensional segment tree.
  --
  -- @since 1.2.3.0
  data SegTree2d s a = SegTree2d
    { -- | The number of nodes.
      --
      -- @since 1.2.3.0
      nSt :: {-# UNPACK #-} !Int,
      -- | The number of distinct \(x\) coordinates.
      --
      -- @since 1.2.3.0
      nxSt :: {-# UNPACK #-} !Int,
      -- | \(\lceil \log_2 (\mathrm{nx} + 1) \rceil\)
      --
      -- @since 1.2.3.0
      logSt :: {-# UNPACK #-} !Int,
      -- | \(2^{\mathrm{logSt}}\)
      --
      -- @since 1.2.3.0
      sizeSt :: {-# UNPACK #-} !Int,
      -- | \(x\) coordinates sorted and uniquified
      --
      -- @since 1.2.3.0
      dictXSt :: {-# UNPACK #-} !(VU.Vector Int),
      -- | \(y\) coordinates sorted (not uniquified)
      --
      -- @since 1.2.3.0
      allYSt :: !(VU.Vector Int),
      -- |
      --
      -- @since 1.2.3.0
      posSt :: !(VU.Vector Int),
      -- |
      --
      -- @since 1.2.3.0
      indptrSt :: !(VU.Vector Int),
      -- | Monoid values
      --
      -- @since 1.2.3.0
      dataSt :: !(VUM.MVector s a),
      -- |
      --
      -- @since 1.2.3.0
      toLeftSt :: !(VU.Vector Int)
    }
  
  -- | \(O(n \log n)\) Creates a `SegTree2d` from a vector of \((x, y)\) coordinates.
  --
  -- @since 1.2.3.0
  {-# INLINE new #-}
  new ::
    (PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | \((x, y)\) vector
    VU.Vector (Int, Int) ->
    -- | Two-dimensional segment tree
    m (SegTree2d (PrimState m) a)
  new xys = stToPrim $ buildST xs ys (VU.replicate n mempty)
    where
      n = VU.length xys
      (!xs, !ys) = VU.unzip xys
  
  -- | \(O(n \log n)\) Creates a `SegTree2d` from vectors of \(x\), \(y\) and \(w\).
  --
  -- @since 1.2.3.0
  {-# INLINE build #-}
  build ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | \(x\) vector
    VU.Vector Int ->
    -- | \(y\) vector
    VU.Vector Int ->
    -- | \(w\) vector
    VU.Vector a ->
    -- | Two-dimensional segment tree
    m (SegTree2d (PrimState m) a)
  build xs ys ws = stToPrim $ buildST xs ys ws
  
  -- | \(O(n \log n)\) Creates a `SegTree2d` from vectors of \((x, y)\) and \(w\).
  --
  -- @since 1.2.3.0
  {-# INLINE build2 #-}
  build2 ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | \((x, y)\) vector
    VU.Vector (Int, Int) ->
    -- | \(w\) vector
    VU.Vector a ->
    -- | Two-dimensional segment tree
    m (SegTree2d (PrimState m) a)
  build2 xys ws = stToPrim $ do
    let (!xs, !ys) = VU.unzip xys
    buildST xs ys ws
  
  -- | \(O(n \log n)\) Creates a `SegTree2d` from a vector of \((x, y, w)\).
  --
  -- @since 1.2.3.0
  {-# INLINE build3 #-}
  build3 :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => VU.Vector (Int, Int, a) -> m (SegTree2d (PrimState m) a)
  build3 xyws = stToPrim $ do
    let (!xs, !ys, !ws) = VU.unzip3 xyws
    buildST xs ys ws
  
  -- -- | \(O(\log n)\) Read the \(k\)-th original point's monoid value.
  -- --
  -- -- @since 1.2.3.0
  -- {-# INLINE read #-}
  -- read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree2d (PrimState m) a -> Int -> m ()
  -- read SegTree2d {..} i = do
  --   let !_ = ACIA.checkIndex "AtCoder.Extra.SegTree2d.read" i nSt
  --   -- FIXME: pos is incorrect
  --   VGM.read dataSt (posSt VG.! i)
  
  -- | \(O(\log n)\) Writes to the \(k\)-th original point's monoid value.
  --
  -- @since 1.2.3.0
  {-# INLINE write #-}
  write ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | Two-dimensional segment tree.
    SegTree2d (PrimState m) a ->
    -- | Original point index.
    Int ->
    -- | New monoid value.
    a ->
    m ()
  write seg i x = stToPrim $ do
    modifyM seg (pure . const x) i
  
  -- | \(O(\log n)\) Given a user function \(f\), modifies the \(k\)-th original point's monoid value
  -- with it.
  --
  -- @since 1.2.3.0
  {-# INLINE modify #-}
  modify ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | Two-dimensional segment tree.
    SegTree2d (PrimState m) a ->
    -- | Function that alters the monoid value.
    (a -> a) ->
    -- | Original point index.
    Int ->
    m ()
  modify seg f i = stToPrim $ do
    modifyM seg (pure . f) i
  
  -- | \(O(\log n)\) Given a user function \(f\), modifies the \(k\)-th original point's monoid value
  -- with it.
  --
  -- @since 1.2.3.0
  {-# INLINEABLE modifyM #-}
  modifyM ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | Two-dimensional segment tree.
    SegTree2d (PrimState m) a ->
    -- | Function that alters the monoid value.
    (a -> m a) ->
    -- | Original point index.
    Int ->
    m ()
  modifyM seg@SegTree2d {..} f rawIdx = do
    let !_ = ACIA.checkIndex "AtCoder.Extra.SegTree2d.modifyM" rawIdx nSt
    inner 1 $ posSt VG.! rawIdx
    where
      inner i p = do
        let indptrI = indptrSt VG.! i
        modifyIST seg f i $ p - indptrI
        when (i < sizeSt) $ do
          let lc = toLeftSt VG.! p - toLeftSt VG.! indptrI
          let rc = (p - indptrI) - lc
          if toLeftSt VG.! (p + 1) - toLeftSt VG.! p /= 0
            then do
              let i' = 2 * i + 0
              let p' = indptrSt VG.! i' + lc
              inner i' p'
            else do
              let i' = 2 * i + 1
              let p' = indptrSt VG.! i' + rc
              inner i' p'
  
  -- | \(O(\log^2 n)\) Returns monoid product \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\).
  --
  -- @since 1.2.3.0
  {-# INLINE prod #-}
  prod ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | Two-dimensional segment tree.
    SegTree2d (PrimState m) a ->
    -- | \(x_1\)
    Int ->
    -- | \(x_2\)
    Int ->
    -- | \(y_1\)
    Int ->
    -- | \(y_2\)
    Int ->
    -- | \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\)
    m a
  prod seg lx rx ly ry = stToPrim $ prodST seg lx rx ly ry
  
  -- | \(O(1)\) Returns monoid product \(\Pi_{p \in [-\infty, \infty) \times [-\infty, \infty)} a_p\).
  --
  -- @since 1.2.3.0
  {-# INLINE allProd #-}
  allProd :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree2d (PrimState m) a -> m a
  allProd seg = fromMaybe mempty <$> VGM.readMaybe (dataSt seg) 1
  
  -- | \(O(\log n)\) Returns the number of points in \([x_1, x_2) \times [y_1, y_2)\).
  --
  -- @since 1.2.3.0
  {-# INLINE count #-}
  count ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | Two-dimensional segment tree.
    SegTree2d (PrimState m) a ->
    -- | \(x_1\)
    Int ->
    -- | \(x_2\)
    Int ->
    -- | \(y_1\)
    Int ->
    -- | \(y_2\)
    Int ->
    -- | The number of points in \([x_1, x_2) \times [y_1, y_2)\).
    m Int
  count seg lx rx ly ry = stToPrim $ countST seg lx rx ly ry
  
  -- -------------------------------------------------------------------------------------------------
  -- Internal
  -- -------------------------------------------------------------------------------------------------
  
  {-# INLINEABLE buildST #-}
  buildST :: forall s a. (HasCallStack, Monoid a, VU.Unbox a) => VU.Vector Int -> VU.Vector Int -> VU.Vector a -> ST s (SegTree2d s a)
  buildST xs ys ws = do
    let nSt = VU.length xs
    let !_ = ACIA.runtimeAssert (nSt == VU.length ys && nSt == VU.length ws) "AtCoder.Extra.SegTree2d.buildST: length mismatch among `xs`, `ys` and `ws`"
  
    -- TODO: use radix sort?
    -- we don't have to take `uniq` though:
    let dictXSt = VU.uniq $ VU.modify VAI.sort xs
    let nxSt = VU.length dictXSt
    let logSt = countTrailingZeros $ ACIB.bitCeil (nxSt + 1)
    let sizeSt = bit logSt
    let compressedXs = VU.map (lowerBound dictXSt) xs
  
    -- TODO: what is this?
    let indptrSt = VU.create $ do
          indptr <- VUM.replicate (2 * sizeSt + 1) (0 :: Int)
          VU.forM_ compressedXs $ \i -> do
            let inner j
                  | j /= 0 = do
                      VGM.modify indptr (+ 1) (j + 1) -- +1 for perfix sum
                      inner (j `div` 2)
                  | otherwise = pure ()
            inner $ i + sizeSt
          -- calculate prefix sum in-place:
          VUM.iforM_ (VUM.init indptr) $ \i x -> do
            VGM.modify indptr (+ x) (i + 1)
          pure indptr
  
    dataSt <- VUM.replicate (2 * VU.last indptrSt) (mempty :: a)
  
    let yis = argsort ys
    let posSt = VU.create $ do
          vec <- VUM.replicate nSt (0 :: Int)
          VU.iforM_ yis $ \i yi -> do
            -- FIXME: It overwrites duplicate yi to i. Isn't it just `accumulate max`?
            VGM.write vec yi i
          pure vec
  
    -- +1 for prefix cumulative sum
    toLeftSt <- do
      toLeft <- VUM.replicate (indptrSt VG.! sizeSt + 1) (0 :: Int)
      ptr <- VU.thaw indptrSt
      VU.forM_ yis $ \rawIdx -> do
        let inner i j
              | i > 0 = do
                  p <- VGM.read ptr i
                  VGM.write ptr i $ p + 1
                  VGM.write dataSt (indptrSt VG.! (i + 1) + p) $ ws VG.! rawIdx
                  when (j /= -1) $ do
                    VGM.write toLeft (p + 1) $ if even j then 1 else 0
                  inner (i `div` 2) i
              | otherwise = pure ()
        let i0 = compressedXs VG.! rawIdx + sizeSt
        inner i0 (-1)
  
      -- calculate prefix cumulative sum in-place:
      VUM.iforM_ (VUM.init toLeft) $ \i x -> do
        VGM.modify toLeft (+ x) (i + 1)
      VU.unsafeFreeze toLeft
  
    for_ [0 .. 2 * sizeSt - 1] $ \i -> do
      let off = 2 * indptrSt VG.! i
      let n = indptrSt VG.! (i + 1) - indptrSt VG.! i
      for_ [n - 1, n - 2 .. 1] $ \j -> do
        xl <- VGM.read dataSt $ off + 2 * j + 0
        xr <- VGM.read dataSt $ off + 2 * j + 1
        VGM.write dataSt (off + j) $! xl <> xr
  
    -- TODO: why not uniquified?
    let allYSt = VU.modify VAI.sort ys
    pure SegTree2d {..}
  
  {-# INLINEABLE modifyIST #-}
  modifyIST :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
  modifyIST SegTree2d {..} f i j0 = do
    -- TODO: LID?
    let lid = indptrSt VG.! i
    let off = 2 * lid
    let inner j_ = do
          when (j_ > 1) $ do
            let j = j_ `div` 2
            xl <- VGM.read dataSt (off + 2 * j + 0)
            xr <- VGM.read dataSt (off + 2 * j + 1)
            VGM.write dataSt (off + j) $! xl <> xr
            inner j
    let j = j0 + indptrSt VG.! (i + 1) - lid
    VGM.modifyM dataSt f (off + j)
    stToPrim $ inner j
  
  {-# INLINEABLE prodST #-}
  prodST :: forall s a. (HasCallStack, Monoid a, VU.Unbox a) => SegTree2d s a -> Int -> Int -> Int -> Int -> ST s a
  prodST seg@SegTree2d {..} lx rx ly ry = do
    let a0 = lowerBound allYSt ly
    let b0 = lowerBound allYSt ry
    dfs mempty 1 0 sizeSt a0 b0
    where
      !_ = ACIA.runtimeAssert (lx <= rx && ly <= ry) "AtCoder.Extra.SegTree2d.prodST: given invalid rectangle"
      !l0 = lowerBound dictXSt lx
      !r0 = lowerBound dictXSt rx
      dfs :: a -> Int -> Int -> Int -> Int -> Int -> ST s a
      dfs !res i l r a b
        -- empty rect
        | a == b = pure res
        -- not intersecting
        | r0 <= l || r <= l0 = pure res
        -- fully contained in [l_0, r_0)
        | l0 <= l && r <= r0 = do
            xi <- prodIST seg i a b
            pure $! xi <> res
        | otherwise = do
            let indptrI = indptrSt VG.! i
            let toLeftI = toLeftSt VG.! indptrI
            let la = toLeftSt VG.! (indptrI + a) - toLeftI
            let ra = a - la
            let lb = toLeftSt VG.! (indptrI + b) - toLeftI
            let rb = b - lb
            let m = (l + r) `div` 2
            res' <- dfs res (2 * i + 0) l m la lb
            dfs res' (2 * i + 1) m r ra rb
  
  {-# INLINEABLE prodIST #-}
  prodIST :: (HasCallStack, Monoid a, VU.Unbox a) => SegTree2d s a -> Int -> Int -> Int -> ST s a
  prodIST SegTree2d {..} i a b = inner mempty (n + a) (n + b - 1)
    where
      lid = indptrSt VG.! i
      off = 2 * lid
      n = indptrSt VG.! (i + 1) - lid
      -- NOTE: we're using inclusive interval [l, r] for simplicity
      inner !res l r
        | l <= r = do
            res' <-
              if testBit l 0
                then (res <>) <$> VGM.read dataSt (off + l)
                else pure res
            res'' <-
              if not $ testBit r 0
                then (<> res') <$> VGM.read dataSt (off + r)
                else pure res'
            inner res'' ((l + 1) .>>. 1) ((r - 1) .>>. 1)
        | otherwise = pure res
  
  {-# INLINEABLE countST #-}
  countST :: forall s a. (HasCallStack, Monoid a, VU.Unbox a) => SegTree2d s a -> Int -> Int -> Int -> Int -> ST s Int
  countST SegTree2d {..} lx rx ly ry = do
    let a0 = lowerBound allYSt ly
    let b0 = lowerBound allYSt ry
    dfs 0 1 0 sizeSt a0 b0
    where
      !_ = ACIA.runtimeAssert (lx <= rx && ly <= ry) "AtCoder.Extra.SegTree2d.countST: given invalid rectangle"
      !l0 = lowerBound dictXSt lx
      !r0 = lowerBound dictXSt rx
      dfs :: Int -> Int -> Int -> Int -> Int -> Int -> ST s Int
      dfs (res :: Int) i l r a b
        -- empty rect
        | a == b = pure res
        -- not intersecting
        | r0 <= l || r <= l0 = pure res
        -- fully contained in [l_0, r_0)
        | l0 <= l && r <= r0 = do
            pure $! res + b - a
        | otherwise = do
            let indptrI = indptrSt VG.! i
            let toLeftI = toLeftSt VG.! indptrI
            let la = toLeftSt VG.! (indptrI + a) - toLeftI
            let ra = a - la
            let lb = toLeftSt VG.! (indptrI + b) - toLeftI
            let rb = b - lb
            let m = (l + r) `div` 2
            res' <- dfs res (2 * i + 0) l m la lb
            dfs res' (2 * i + 1) m r ra rb
  