content: {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ViewPatterns #-}
  
  -- | Same test with non-persistent version
  module Tests.Extra.DynSegTree.Persistent where
  
  import AtCoder.Extra.DynSegTree.Persistent qualified as Seg
  import AtCoder.ModInt qualified as M
  import Control.Monad.ST (RealWorld, runST)
  import Data.Foldable (for_)
  import Data.Semigroup (Sum (..))
  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 Test.QuickCheck.Monadic as QCM
  import Test.Tasty
  import Test.Tasty.QuickCheck as QC
  import Tests.Util
  
  type Mint = M.ModInt998244353
  
  modInt :: Int -> Mint
  modInt = M.new
  
  data Init = Init
    { n :: {-# UNPACK #-} !Int,
      q :: {-# UNPACK #-} !Int,
      refM :: !(IO (VUM.MVector RealWorld (Sum Mint))),
      segM :: !(IO (Seg.DynSegTree RealWorld (Sum Mint)))
    }
  
  instance Show Init where
    show Init {..} = show n
  
  initialProd :: Int -> Int -> Sum Mint
  initialProd l r = Sum . modInt $ sum [l .. r - 1]
  
  instance QC.Arbitrary Init where
    arbitrary = do
      n <- QC.chooseInt (1, 256)
      q <- QC.chooseInt (1, 5 * n)
      l0 <- QC.chooseInt (-256, 256)
      -- initial node values: l0, l0 + 1, l0 + 2, ..
      let cap = Seg.recommendedCapacity n q
      pure
        . Init n q (VUM.generate n (\i -> initialProd (i + l0) (i + l0 + 1)))
        $ Seg.buildWith cap l0 (l0 + n) initialProd
  
  data Query
    = Write !Int !(Sum Mint)
    | Modify !Int !(Sum Mint)
    | ModifyM !Int !(Sum Mint)
    | Prod !(Int, Int)
    | AllProd
    | ResetInterval !(Int, Int)
    deriving (Show)
  
  -- | Arbitrary return type for the `Query` result.
  data Result
    = None
    | S !(Sum Mint)
    | F !(VU.Vector (Sum Mint))
    deriving (Show, Eq)
  
  queryGen :: Int -> QC.Gen Query
  queryGen n = do
    QC.oneof
      [ Write <$> k <*> v,
        Modify <$> k <*> v,
        ModifyM <$> k <*> v,
        Prod <$> i,
        pure AllProd,
        ResetInterval <$> i
      ]
    where
      k = QC.chooseInt (0, n - 1)
      v = Sum . modInt <$> QC.arbitrary
      i = intervalGen n
  
  -- | containers. (referential implementation)
  handleRef :: Int -> VUM.MVector RealWorld (Sum Mint) -> Query -> IO Result
  handleRef l0 vec q = case q of
    Write k v -> do
      VGM.write vec k v
      pure None
    Modify k v -> do
      VGM.modify vec (+ v) k
      pure None
    ModifyM k v -> do
      VGM.modify vec (+ v) k
      pure None
    Prod (!l, !r) -> do
      (S <$>) . VGM.foldl' (<>) mempty $ VGM.slice l (r - l) vec
    AllProd -> do
      S <$> VGM.foldl' (<>) mempty vec
    ResetInterval (!l, !r) -> do
      for_ [l .. r - 1] $ \i -> do
        VGM.write vec i $! initialProd (l0 + i) (l0 + i + 1)
      pure None
  
  -- | ac-library-hs.
  handleAcl :: (HasCallStack) => Seg.DynSegTree RealWorld (Sum Mint) -> VUM.MVector RealWorld Seg.Index -> Query -> IO Result
  handleAcl seg rootVec q = do
    root <- VGM.read rootVec 0
    case q of
      Write (d -> k) v -> do
        VGM.write rootVec 0 =<< Seg.write seg root k v
        VGM.write rootVec 0 =<< Seg.write seg root k v
        pure None
      Modify (d -> k) v -> do
        VGM.write rootVec 0 =<< Seg.modify seg root (+ v) k
        pure None
      ModifyM (d -> k) v -> do
        VGM.write rootVec 0 =<< Seg.modifyM seg root (pure . (+ v)) k
        pure None
      Prod (d -> l, d -> r) -> do
        S <$> Seg.prod seg root l r
      AllProd -> do
        S <$> Seg.allProd seg root
      ResetInterval (d -> l, d -> r) -> do
        VGM.write rootVec 0 =<< Seg.resetInterval seg root l r
        pure None
    where
      d = (+ Seg.l0Dst seg)
  
  prop_randomTest :: Init -> QC.Property
  prop_randomTest Init {..} = QCM.monadicIO $ do
    qs <- QCM.pick $ QC.vectorOf q (queryGen n)
    ref <- QCM.run refM
    seg <- QCM.run segM
    let l0 = Seg.l0Dst seg
    let bounds = (Seg.l0Dst seg, Seg.r0Dst seg)
    root <- QCM.run $ Seg.newRoot seg
    rootVec <- QCM.run $ VUM.replicate 1 root
    for_ qs $ \query -> do
      expected <- QCM.run $ handleRef l0 ref query
      res <- QCM.run $ handleAcl seg rootVec query
      QCM.assertWith (expected == res) $ show (bounds, query, expected, res)
  
  -- prop_foldl is tested with large array verification problem
  
  prop_maxRight :: Int -> QC.NonEmptyList (QC.NonNegative Int) -> QC.Property
  prop_maxRight xRef (QC.NonEmpty xs_) =
    let xs = VU.fromList $ map (\(QC.NonNegative x) -> x) xs_
        expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs
        res = runST $ do
          seg <- Seg.new @_ @(Sum Int) (2 * VU.length xs) 0 (VU.length xs)
          root <- Seg.newSeq seg $ VU.map Sum xs
          Seg.maxRight seg root (<= Sum xRef)
     in expected QC.=== res
  
  tests :: [TestTree]
  tests =
    [ QC.testProperty "randomTest" prop_randomTest,
      QC.testProperty "maxRight" prop_maxRight
    ]
  