content: {-# LANGUAGE RecordWildCards #-}
  
  -- | Static, \(k\)-dimensional tree \((k = 2)\) with lazily propagated monoid actions and
  -- commutative acted monoids.
  --
  -- - Point coordinates are fixed on `build` and cannot be moved or added later.
  -- - Multiple points can exist at the same coordinate.
  --
  -- ==== __Examples__
  -- >>> import AtCoder.Extra.LazyKdTree qualified as LKT
  -- >>> import AtCoder.Extra.Monoid.Affine1 (Affine1)
  -- >>> import AtCoder.Extra.Monoid.Affine1 qualified as Affine1
  -- >>> import Data.Semigroup (Sum (..))
  -- >>> import Data.Vector.Unboxed qualified as VU
  -- >>> let xyws = VU.fromList [(0, 0, Sum 1), (1, 1, Sum 2), (4, 2, Sum 3)]
  -- >>> lkt <- LKT.build3 @_ @(Affine1 Int) @(Sum Int) xyws
  --
  -- >>> -- Get monoid product in [0, 2) x [0, 2)
  -- >>> LKT.prod lkt 0 2 0 2
  -- Sum {getSum = 3}
  --
  -- >>> LKT.applyIn lkt 0 2 0 2 $ Affine1.new 2 1
  -- >>> LKT.prod lkt 0 2 0 2
  -- Sum {getSum = 8}
  --
  -- >>> LKT.write lkt 0 $ Sum 10
  -- >>> LKT.prod lkt 0 2 0 2
  -- Sum {getSum = 15}
  --
  -- @since 1.2.2.0
  module AtCoder.Extra.LazyKdTree
    ( -- * K-dimensional tree
      LazyKdTree (..),
  
      -- * Re-exports
      SegAct (..),
  
      -- * Constructors
      new,
      build,
      build2,
      build3,
  
      -- * Write
      write,
      modify,
      modifyM,
  
      -- * Monoid products
      prod,
      allProd,
  
      -- * Apply
      applyIn,
    )
  where
  
  import AtCoder.Internal.Assert qualified as ACIA
  import AtCoder.Internal.Bit qualified as ACIB
  import AtCoder.LazySegTree (SegAct (..))
  import Control.Monad (unless, 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.Ord (comparing)
  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)
  
  -- | Static, \(k\)-dimensional tree \((k = 2)\) with lazily propagated monoid actions and
  -- commutative monoids.
  --
  -- @since 1.2.2.0
  data LazyKdTree s f a = LazyKdTree
    { -- | The number of points in the \(k\)-d tree.
      --
      -- @since 1.2.2.0
      nLkt :: {-# UNPACK #-} !Int,
      -- | \(\lceil \log_2 (n + 1) \rceil\)
      --
      -- @since 1.2.2.0
      logLkt :: {-# UNPACK #-} !Int,
      -- | Rectangle information: inclusive (closed) ranges \([x_1, x_2) \times [y_1, y_2)\).
      --
      -- @since 1.2.2.0
      incRectsLkt :: !(VU.Vector (Int, Int, Int, Int)),
      -- | Rectangle information: monoid values.
      --
      -- @since 1.2.2.0
      dataLkt :: !(VUM.MVector s a),
      -- | Rectangle information: lazily propagated monoid actions for children.
      --
      -- @since 1.2.2.0
      lazyLkt :: !(VUM.MVector s f),
      -- | Rectangle information: the number of vertices in the rectangle.
      --
      -- @since 1.2.2.0
      sizeLkt :: !(VU.Vector Int),
      -- | Maps original vertices into the belonging rectangle index.
      --
      -- @since 1.2.2.0
      posLkt :: !(VU.Vector Int)
    }
  
  -- | \(O(n \log n)\) Creates a `LazyKdTree` from @xs@ and @ys@.
  --
  -- ==== Constraints
  -- - \(|\mathrm{xs}| = |\mathrm{ys}|\)
  --
  -- @since 1.2.3.0
  {-# INLINE new #-}
  new ::
    (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
    -- | \(x\) coordinates.
    VU.Vector Int ->
    -- | \(y\) coordinates.
    VU.Vector Int ->
    -- | `LazyKdTree`.
    m (LazyKdTree (PrimState m) f a)
  new xs ys = stToPrim $ buildST xs ys (VU.replicate (VU.length xs) mempty)
  
  -- | \(O(n \log n)\) Creates a `LazyKdTree` from @xs@, @ys@ and @ws@ vectors.
  --
  -- ==== Constraints
  -- - \(|\mathrm{xs}| = |\mathrm{ys}| = |\mathrm{vs}|\).
  --
  -- @since 1.2.2.0
  {-# INLINE build #-}
  build ::
    (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
    -- | \(x\) coordinates.
    VU.Vector Int ->
    -- | \(y\) coordinates.
    VU.Vector Int ->
    -- | monoid \(v\)alues.
    VU.Vector a ->
    -- | `LazyKdTree`.
    m (LazyKdTree (PrimState m) f a)
  build xs ys vs = stToPrim $ buildST xs ys vs
  
  -- | \(O(n \log n)\) Creates a `LazyKdTree` from @xys@ and @ws@ vectors.
  --
  -- ==== Constraints
  -- - \(|\mathrm{xys}| = |\mathrm{vs}|\).
  --
  -- @since 1.2.2.0
  {-# INLINE build2 #-}
  build2 ::
    (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
    -- | \((x, y)\) coordinates.
    VU.Vector (Int, Int) ->
    -- | Monoid \(v\)alues.
    VU.Vector a ->
    -- | `LazyKdTree`.
    m (LazyKdTree (PrimState m) f a)
  build2 xys ws = stToPrim $ buildST xs ys ws
    where
      (!xs, !ys) = VU.unzip xys
  
  -- | \(O(n \log n)\) Creates a `LazyKdTree` from a @xyws@ vector.
  --
  -- @since 1.2.2.0
  {-# INLINE build3 #-}
  build3 ::
    (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
    -- | \((x, y, v)\) tuples.
    VU.Vector (Int, Int, a) ->
    -- | `LazyKdTree`.
    m (LazyKdTree (PrimState m) f a)
  build3 xyws = stToPrim $ buildST xs ys ws
    where
      (!xs, !ys, !ws) = VU.unzip3 xyws
  
  -- | \(O(\log n)\) Writes to the \(k\)-th point's monoid value.
  --
  -- @since 1.2.2.0
  {-# INLINE write #-}
  write ::
    (HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
    -- | `LazyKdTree`.
    LazyKdTree (PrimState m) f a ->
    -- | Original vertex index.
    Int ->
    -- | Monoid value.
    a ->
    -- | Monadic tuple.
    m ()
  write kt i x = stToPrim $ modifyM kt (pure . const x) i
  
  -- | \(O(\log n)\) Given a user function \(f\), modifies the \(k\)-th point's monoid value with it.
  --
  -- @since 1.2.2.0
  {-# INLINE modify #-}
  modify ::
    (HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
    -- | `LazyKdTree`.
    LazyKdTree (PrimState m) f a ->
    -- | \(f\): Creates a new monoid value from the old one.
    (a -> a) ->
    -- | \(k\): Original vertex index.
    Int ->
    -- | Monadic tuple.
    m ()
  modify kt f i = stToPrim $ modifyM kt (pure . f) i
  
  -- | \(O(\log n)\) Given a user function \(f\), modifies the \(k\)-th point's monoid value with it.
  --
  -- @since 1.2.2.0
  {-# INLINEABLE modifyM #-}
  modifyM ::
    (HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
    -- | `LazyKdTree`.
    LazyKdTree (PrimState m) f a ->
    -- | \(f\): Creates a new monoid value from the old one.
    (a -> m a) ->
    -- | \(k\): Original vertex index.
    Int ->
    -- | Monadic tuple.
    m ()
  modifyM kt@LazyKdTree {..} f i0 = do
    let !_ = ACIA.checkIndex "AtCoder.Extra.LazyKdTree.modifyM" i0 nLkt
    let i_ = posLkt VG.! i0
    -- propagate lazily propagated monoid actions from the root:
    stToPrim $ for_ [logLkt, logLkt - 1 .. 1] $ \k -> do
      pushST kt (i_ .>>. k)
    VGM.modifyM dataLkt f i_
    -- update parents:
    let inner i
          | i <= 1 = pure ()
          | otherwise = do
              let i' = i `div` 2
              xl <- VGM.read dataLkt (2 * i' + 0)
              xr <- VGM.read dataLkt (2 * i' + 1)
              VGM.write dataLkt i' $! xl <> xr
              inner i'
    stToPrim $ inner i_
  
  -- | \(O(\sqrt n)\) Returns monoid product \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\).
  --
  -- @since 1.2.2.0
  {-# INLINE prod #-}
  prod ::
    (HasCallStack, PrimMonad m, Eq f, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) =>
    -- | `LazyKdTree`
    LazyKdTree (PrimState m) f 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 kt x1 x2 y1 y2 = stToPrim $ prodST kt x1 x2 y1 y2
  
  -- | \(O(1)\) Returns monoid product \(\Pi_{p \in [-\infty, \infty) \times [-\infty, \infty)} a_p\).
  --
  -- @since 1.2.2.0
  {-# INLINE allProd #-}
  allProd ::
    (PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | `LazyKdTree`
    LazyKdTree (PrimState m) f a ->
    -- | \(\Pi_{p \in [-\infty, \infty) \times [-\infty, \infty)} a_p\)
    m a
  allProd kt = do
    -- In case of zero vertices, use `Maybe`:
    fromMaybe mempty <$> VGM.readMaybe (dataLkt kt) 1
  
  -- | \(O(\log n)\) Given a rectangle \([x_1, x_2) \times [y_1, y_2)\), applies a monoid action to
  -- it.
  --
  -- @since 1.2.2.0
  {-# INLINE applyIn #-}
  applyIn ::
    (HasCallStack, PrimMonad m, Eq f, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) =>
    -- | `LazyKdTree`
    LazyKdTree (PrimState m) f a ->
    -- | \(x_1\)
    Int ->
    -- | \(x_2\)
    Int ->
    -- | \(y_1\)
    Int ->
    -- | \(y_2\)
    Int ->
    -- | \(f\)
    f ->
    -- | Monadic tuple.
    m ()
  applyIn kt x1 x2 y1 y2 f = stToPrim $ applyInST kt 1 x1 x2 y1 y2 f
  
  -- -------------------------------------------------------------------------------------------------
  -- Private
  -- -------------------------------------------------------------------------------------------------
  
  {-# INLINEABLE buildST #-}
  buildST :: forall s f a. (HasCallStack, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) => VU.Vector Int -> VU.Vector Int -> VU.Vector a -> ST s (LazyKdTree s f a)
  buildST xs0 ys0 vs0 = do
    let nLkt = VU.length xs0
    let !_ = ACIA.runtimeAssert (nLkt == VU.length ys0 && nLkt == VU.length vs0) "AtCoder.Extra.LazyKdTree.buildST: the length of `xs`, `ys` and `vs` must be equal"
    if nLkt == 0
      then do
        let logLkt = 0
        dataLkt <- VUM.new 0
        lazyLkt <- VUM.new 0
        let incRectsLkt = VU.empty
        let sizeLkt = VU.empty
        let posLkt = VU.empty
        pure LazyKdTree {..}
      else do
        let logLkt = countTrailingZeros $ ACIB.bitCeil (nLkt + 1)
        dataLkt <- VUM.unsafeNew (bit (logLkt + 1))
        lazyLkt <- VUM.replicate (bit logLkt) mempty
        incRectsVec <- VUM.replicate (bit (logLkt + 1)) (maxBound, minBound, maxBound, minBound)
        size <- VUM.unsafeNew (bit (logLkt + 1))
        pos <- VUM.unsafeNew nLkt
        let VUM.MV_4 _ xMins xMaxes yMins yMaxes = incRectsVec
  
        -- - idx: rectangle index (one-based)
        -- - xs, ys, vs: point information (x, y and monoid value)
        -- - ids: maps sorted vertices to the original vertex indices
        -- - divX: represents hyperplane direction for point partition
        let buildSubtree :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector a -> VU.Vector Int -> Bool -> ST s ()
            buildSubtree idx xs ys vs ids divX = do
              let n = VU.length xs
              VGM.write size idx n
  
              -- retrieve the bounds:
              let (!xMin, !xMax, !yMin, !yMax) =
                    VU.foldl'
                      (\(!a, !b, !c, !d) (!x, !y) -> (min a x, max b x, min c y, max d y))
                      (maxBound, minBound, maxBound, minBound)
                      $ VU.zip xs ys
              VGM.modify xMins (min xMin) idx
              VGM.modify xMaxes (max xMax) idx
              VGM.modify yMins (min yMin) idx
              VGM.modify yMaxes (max yMax) idx
  
              if n == 1
                then do
                  -- it's a terminal. note that it's not always a leaf; the case is handled carefully in
                  -- other methods
                  VGM.write dataLkt idx $ vs VG.! 0
                  -- record original vertex index -> rectangle index
                  VGM.write pos (ids VG.! 0) idx
                else do
                  -- partition the vertices into two:
                  let m = n `div` 2
                  let is = VU.create $ do
                        vec <- VUM.generate n id
                        if divX
                          then VAI.selectBy (comparing (xs VG.!)) vec m
                          else VAI.selectBy (comparing (ys VG.!)) vec m
                        pure vec
  
                  -- TODO: permute in-place?
                  let (!xsL, !xsR) = VG.splitAt m $ VG.backpermute xs is
                  let (!ysL, !ysR) = VG.splitAt m $ VG.backpermute ys is
                  let (!vsL, !vsR) = VG.splitAt m $ VG.backpermute vs is
                  let (!idsL, !idsR) = VG.splitAt m $ VG.backpermute ids is
  
                  -- build the subtree:
                  buildSubtree (2 * idx + 0) xsL ysL vsL idsL (not divX)
                  buildSubtree (2 * idx + 1) xsR ysR vsR idsR (not divX)
                  xl <- VGM.read dataLkt (2 * idx + 0)
                  xr <- VGM.read dataLkt (2 * idx + 1)
                  VGM.write dataLkt idx $! xl <> xr
  
        buildSubtree 1 xs0 ys0 vs0 (VU.generate nLkt id) True
        sizeLkt <- VU.unsafeFreeze size
        posLkt <- VU.unsafeFreeze pos
        incRectsLkt <- VU.unsafeFreeze incRectsVec
        pure LazyKdTree {..}
  
  {-# INLINE applyAtST #-}
  applyAtST :: (SegAct f a, VU.Unbox f, VU.Unbox a) => LazyKdTree s f a -> Int -> f -> ST s ()
  applyAtST LazyKdTree {..} i f = do
    -- NOTE: Here we're asssuming each monoid value has length one. If you need a monoid of length
    -- zero, e.g., if you're just reserving new point insertion, you must not rely on
    -- `segActWithLength`. You might want to use `V2` instead of `Sum`.
    let len = sizeLkt VG.! i
    VGM.modify dataLkt (segActWithLength len f) i
    when (i < bit logLkt) $ do
      VGM.modify lazyLkt (f <>) i
  
  -- TODO: consider `INLINE`?
  {-# INLINE pushST #-}
  pushST :: (SegAct f a, Eq f, VU.Unbox f, VU.Unbox a) => LazyKdTree s f a -> Int -> ST s ()
  pushST kt@LazyKdTree {..} i = do
    lazy <- VGM.read lazyLkt i
    unless (lazy == mempty) $ do
      applyAtST kt (2 * i + 0) lazy
      applyAtST kt (2 * i + 1) lazy
      VGM.write lazyLkt i mempty
  
  {-# INLINEABLE prodST #-}
  prodST :: (HasCallStack, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => LazyKdTree s f a -> Int -> Int -> Int -> Int -> ST s a
  prodST kt@LazyKdTree {..} x1 x2 y1 y2
    | x1 >= x2 || y1 >= y2 = pure mempty
    | otherwise = inner 1
    where
      inner i = case incRectsLkt VG.!? i of
        Nothing -> pure mempty
        Just (!xl, !xr, !yl, !yr)
          -- TODO: what is this?
          | xl > xr -> pure mempty
          -- not intersecting
          | x2 <= xl || x1 > xr || y2 <= yl || y1 > yr -> pure mempty
          -- the rectangle is fully contained by the query:
          | x1 <= xl && xr < x2 && y1 <= yl && yr < y2 -> do
              VGM.read dataLkt i
          | otherwise -> do
              pushST kt i
              l <- inner (2 * i + 0)
              r <- inner (2 * i + 1)
              pure $! l <> r
  
  {-# INLINEABLE applyInST #-}
  applyInST :: (HasCallStack, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => LazyKdTree s f a -> Int -> Int -> Int -> Int -> Int -> f -> ST s ()
  applyInST kt@LazyKdTree {..} i0 x1 x2 y1 y2 f
    | x1 >= x2 || y1 >= y2 = pure ()
    | otherwise = inner i0
    where
      inner i = case incRectsLkt VG.!? i of
        Nothing -> pure mempty
        Just (!xl, !xr, !yl, !yr)
          | xl > xr -> pure ()
          -- not intersecting
          | x2 <= xl || x1 > xr || y2 <= yl || y1 > yr -> pure ()
          -- the rectangle is fully contained by the query:
          | x1 <= xl && xr < x2 && y1 <= yl && yr < y2 -> do
              applyAtST kt i f
          | otherwise -> do
              pushST kt i
              inner (2 * i + 0)
              inner (2 * i + 1)
              l <- VGM.read dataLkt (2 * i + 0)
              r <- VGM.read dataLkt (2 * i + 1)
              VGM.write dataLkt i $! l <> r
  