module Data.FenwickTree.RangeAdd where

import Control.Monad
import Control.Monad.Primitive
import Data.Bits
import Data.Function
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

import My.Prelude (floorPowerOf2, rep1, rev1)

newtype RangeAddFenwickTree s a = RangeAddFenwickTree (UM.MVector s a)

newRangeAddFenwickTree ::
  (Num a, U.Unbox a, PrimMonad m) =>
  Int ->
  m (RangeAddFenwickTree (PrimState m) a)
newRangeAddFenwickTree :: forall a (m :: * -> *).
(Num a, Unbox a, PrimMonad m) =>
Int -> m (RangeAddFenwickTree (PrimState m) a)
newRangeAddFenwickTree Int
n = MVector (PrimState m) a -> RangeAddFenwickTree (PrimState m) a
forall s a. MVector s a -> RangeAddFenwickTree s a
RangeAddFenwickTree (MVector (PrimState m) a -> RangeAddFenwickTree (PrimState m) a)
-> m (MVector (PrimState m) a)
-> m (RangeAddFenwickTree (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UM.replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
0
{-# INLINE newRangeAddFenwickTree #-}

buildRangeAddFenwickTree ::
  forall a m.
  (Num a, U.Unbox a, PrimMonad m) =>
  U.Vector a ->
  m (RangeAddFenwickTree (PrimState m) a)
buildRangeAddFenwickTree :: forall a (m :: * -> *).
(Num a, Unbox a, PrimMonad m) =>
Vector a -> m (RangeAddFenwickTree (PrimState m) a)
buildRangeAddFenwickTree Vector a
xs = do
  let n :: Int
n = Vector a -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector a
xs
  ft <- Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
UM.unsafeNew (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  UM.write ft 0 0
  U.unsafeCopy (UM.tail ft) xs
  rev1 n $ \Int
i -> do
    x <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
ft (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    UM.unsafeModify ft (subtract x) i
  rep1 n $ \Int
i -> do
    let j :: Int
j = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
i))
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      fti <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
ft Int
i
      UM.unsafeModify ft (+ fti) j
  return $ RangeAddFenwickTree ft
{-# INLINE buildRangeAddFenwickTree #-}

addFromTo ::
  (Num a, U.Unbox a, PrimMonad m) =>
  RangeAddFenwickTree (PrimState m) a ->
  Int ->
  Int ->
  a ->
  m ()
addFromTo :: forall a (m :: * -> *).
(Num a, Unbox a, PrimMonad m) =>
RangeAddFenwickTree (PrimState m) a -> Int -> Int -> a -> m ()
addFromTo (RangeAddFenwickTree MVector (PrimState m) a
ft) Int
l Int
r a
x = do
  Int -> a -> m ()
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
Int -> a -> f ()
add Int
l a
x
  Int -> a -> m ()
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
Int -> a -> f ()
add Int
r (-a
x)
  where
    !n :: Int
n = MVector (PrimState m) a -> Int
forall a s. Unbox a => MVector s a -> Int
UM.length MVector (PrimState m) a
ft
    add :: Int -> a -> f ()
add Int
k a
v = (((Int -> f ()) -> Int -> f ()) -> Int -> f ())
-> Int -> ((Int -> f ()) -> Int -> f ()) -> f ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> f ()) -> Int -> f ()) -> Int -> f ()
forall a. (a -> a) -> a
fix (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (((Int -> f ()) -> Int -> f ()) -> f ())
-> ((Int -> f ()) -> Int -> f ()) -> f ()
forall a b. (a -> b) -> a -> b
$ \Int -> f ()
loop !Int
i -> do
      Bool -> f () -> f ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n) (f () -> f ()) -> f () -> f ()
forall a b. (a -> b) -> a -> b
$ do
        MVector (PrimState f) a -> (a -> a) -> Int -> f ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
UM.unsafeModify MVector (PrimState m) a
MVector (PrimState f) a
ft (a -> a -> a
forall a. Num a => a -> a -> a
+ a
v) Int
i
        Int -> f ()
loop (Int -> f ()) -> Int -> f ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
i))
{-# INLINE addFromTo #-}

-- | /O(log n)/
readRAFT ::
  (Num a, U.Unbox a, PrimMonad m) =>
  RangeAddFenwickTree (PrimState m) a ->
  Int ->
  m a
readRAFT :: forall a (m :: * -> *).
(Num a, Unbox a, PrimMonad m) =>
RangeAddFenwickTree (PrimState m) a -> Int -> m a
readRAFT (RangeAddFenwickTree MVector (PrimState m) a
ft) Int
k = a -> Int -> m a
forall {m :: * -> *}.
(PrimState m ~ PrimState m, PrimMonad m) =>
a -> Int -> m a
go a
0 (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  where
    go :: a -> Int -> m a
go !a
acc !Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = do
          xi <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
MVector (PrimState m) a
ft Int
i
          go (acc + xi) (i - (i .&. (-i)))
      | Bool
otherwise = a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
acc
{-# INLINE readRAFT #-}

-- | /O(log n)/
writeRAFT ::
  (Num a, U.Unbox a, PrimMonad m) =>
  RangeAddFenwickTree (PrimState m) a ->
  Int ->
  a ->
  m ()
writeRAFT :: forall a (m :: * -> *).
(Num a, Unbox a, PrimMonad m) =>
RangeAddFenwickTree (PrimState m) a -> Int -> a -> m ()
writeRAFT RangeAddFenwickTree (PrimState m) a
ft Int
i a
x = do
  fi <- RangeAddFenwickTree (PrimState m) a -> Int -> m a
forall a (m :: * -> *).
(Num a, Unbox a, PrimMonad m) =>
RangeAddFenwickTree (PrimState m) a -> Int -> m a
readRAFT RangeAddFenwickTree (PrimState m) a
ft Int
i
  addFromTo ft i (i + 1) (x - fi)
{-# INLINE writeRAFT #-}

{- |
>>> ft <- buildRangeAddFenwickTree @Int (U.fromList [1,2,2,3,3,3])
>>> lowerBoundRAFT ft 2
1
>>> lowerBoundRAFT ft 0
0
>>> lowerBoundRAFT ft 10
6
-}
lowerBoundRAFT ::
  (U.Unbox a, Num a, Ord a, PrimMonad m) =>
  RangeAddFenwickTree (PrimState m) a ->
  a ->
  m Int
lowerBoundRAFT :: forall a (m :: * -> *).
(Unbox a, Num a, Ord a, PrimMonad m) =>
RangeAddFenwickTree (PrimState m) a -> a -> m Int
lowerBoundRAFT (RangeAddFenwickTree MVector (PrimState m) a
ft) a
s0
  | a
s0 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0 = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0
  | Bool
otherwise = a -> Int -> Int -> m Int
forall {m :: * -> *}.
(PrimState m ~ PrimState m, PrimMonad m) =>
a -> Int -> Int -> m Int
go a
s0 (Int -> Int
floorPowerOf2 Int
n) Int
0
  where
    !n :: Int
n = MVector (PrimState m) a -> Int
forall a s. Unbox a => MVector s a -> Int
UM.length MVector (PrimState m) a
ft
    go :: a -> Int -> Int -> m Int
go !a
s !Int
w !Int
i
      | Int
w Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i
      | Bool
otherwise = do
          if Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
            then do
              fiw <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
MVector (PrimState m) a
ft (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w)
              if fiw < s
                then go (s - fiw) (w !>>. 1) (i + w)
                else go s (w !>>. 1) i
            else a -> Int -> Int -> m Int
go a
s (Int
w Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!>>. Int
1) Int
i
{-# INLINE lowerBoundRAFT #-}

{- |
>>> ft <- buildRangeAddFenwickTree @Int (U.fromList [1,2,2,3,3,3])
>>> upperBoundRAFT ft 2
3
>>> upperBoundRAFT ft 0
0
>>> upperBoundRAFT ft 10
6
-}
upperBoundRAFT ::
  (U.Unbox a, Num a, Ord a, PrimMonad m) =>
  RangeAddFenwickTree (PrimState m) a ->
  a ->
  m Int
upperBoundRAFT :: forall a (m :: * -> *).
(Unbox a, Num a, Ord a, PrimMonad m) =>
RangeAddFenwickTree (PrimState m) a -> a -> m Int
upperBoundRAFT (RangeAddFenwickTree MVector (PrimState m) a
ft) a
s0
  | a
s0 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0
  | Bool
otherwise = a -> Int -> Int -> m Int
forall {m :: * -> *}.
(PrimState m ~ PrimState m, PrimMonad m) =>
a -> Int -> Int -> m Int
go a
s0 (Int -> Int
floorPowerOf2 Int
n) Int
0
  where
    !n :: Int
n = MVector (PrimState m) a -> Int
forall a s. Unbox a => MVector s a -> Int
UM.length MVector (PrimState m) a
ft
    go :: a -> Int -> Int -> m Int
go !a
s !Int
w !Int
i
      | Int
w Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i
      | Bool
otherwise = do
          if Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
            then do
              fiw <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
MVector (PrimState m) a
ft (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w)
              if fiw <= s
                then go (s - fiw) (w !>>. 1) (i + w)
                else go s (w !>>. 1) i
            else a -> Int -> Int -> m Int
go a
s (Int
w Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!>>. Int
1) Int
i
{-# INLINE upperBoundRAFT #-}