content: {-# LANGUAGE DerivingVia #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE TypeFamilies #-}
  
  module Tests.LazySegTreeStress (tests) where
  
  import AtCoder.LazySegTree qualified as LST
  import Control.Monad.Primitive (PrimMonad, PrimState)
  import Control.Monad.State (MonadIO, liftIO)
  import Control.Monad.State.Strict (StateT (..))
  import Data.Foldable
  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 System.Random.Stateful
  import Test.Tasty
  import Test.Tasty.HUnit
  
  -- | Naive implementation of lazy segment tree.
  newtype TimeManager s = TimeManager (VUM.MVector s Int)
  
  action :: (PrimMonad m) => TimeManager (PrimState m) -> Int -> Int -> Int -> m ()
  action (TimeManager vec) l r newTime = do
    for_ [l .. r - 1] $ \i -> do
      VUM.write vec i newTime
  
  prod :: (PrimMonad m) => TimeManager (PrimState m) -> Int -> Int -> m Int
  prod (TimeManager vec) l r =
    VU.foldM'
      ( \ !acc i -> do
          x <- VUM.read vec i
          pure $! max acc x
      )
      (-1 :: Int)
      (VU.generate (r - l) (+ l))
  
  freezeTM :: (PrimMonad m) => TimeManager (PrimState m) -> m (VU.Vector Int)
  freezeTM (TimeManager vec) = VU.freeze vec
  
  -- | S (l, r, time): half-open interval with generation.
  newtype S = S SRepr
    deriving (Eq, Ord, Show)
  
  type SRepr = (Int, Int, Int)
  
  instance Semigroup S where
    l@(S (!l1, !_, !time1)) <> r@(S (!l2, !r2, !time2))
      | l1 == -1 = r
      | l2 == -1 = l
      | otherwise = S (l1, r2, max time1 time2)
  
  instance Monoid S where
    mempty = S (-1, -1, -1)
  
  newtype instance VU.MVector s S = MV_S (VU.MVector s SRepr)
  
  newtype instance VU.Vector S = V_S (VU.Vector SRepr)
  
  deriving instance VGM.MVector VUM.MVector S
  
  deriving instance VG.Vector VU.Vector S
  
  instance VU.Unbox S
  
  -- | T newTime: generation.
  newtype T = T TRepr
    deriving (Eq, Ord, Show)
  
  type TRepr = Int
  
  instance Semigroup T where
    -- new <> old
    (T t1) <> r@(T t2)
      | t1 == -1 = r
      | t1 <= t2 = error "unreachable"
      | otherwise = T t1
  
  instance Monoid T where
    mempty = T (-1)
  
  instance LST.SegAct T S where
    segAct t@(T newTime) s@(S (!l, !r, !time))
      | t == mempty = s
      | time >= newTime = error "unreachable"
      | otherwise = S (l, r, newTime)
  
  newtype instance VU.MVector s T = MV_T (VU.MVector s TRepr)
  
  newtype instance VU.Vector T = V_T (VU.Vector TRepr)
  
  deriving instance VGM.MVector VUM.MVector T
  
  deriving instance VG.Vector VU.Vector T
  
  instance VU.Unbox T
  
  -- | Random choice of 2 disjoint elements from [lower, upper] (l < r)
  uniformPairM :: (StatefulGen g m) => (Int, Int) -> g -> m (Int, Int)
  uniformPairM rng@(!lower, !upper) g
    | upper - lower < 1 = error "invalid argument"
    | otherwise = do
        a <- uniformRM rng g
        b <- uniformRM rng g
        if a == b
          then uniformPairM rng g -- retry
          else pure (min a b, max a b)
  
  -- TODO: write type?
  -- testDriver :: Int ->
  --    (Int -> LST.LazySegTree (PrimState (StateT UniformGen IO )
  testDriver ::
    (MonadIO m, PrimMonad m) =>
    (Int, Int) ->
    Int ->
    (Int -> LST.LazySegTree (PrimState m) T S -> TimeManager (PrimState m) -> Int -> Int -> Int -> StateT StdGen m Int) ->
    m ()
  testDriver tyRange@(!_, !_) nRepeat f = do
    for_ [1 .. 30] $ \n -> do
      for_ [1 .. 10 - 1] $ \_ -> do
        seg0 <- LST.build $ VU.generate n $ \i -> S (i, i + 1, -1)
        tm <- TimeManager <$> VUM.replicate n (-1)
        runStateGenT_ (mkStdGen 42) $ \g -> do
          VU.foldM'_
            ( \now _ -> do
                ty :: Int <- uniformRM tyRange g
                (!l, !r) <- uniformPairM (0, n) g
                f now seg0 tm ty l r
            )
            (0 :: Int)
            (VU.generate nRepeat id)
  
        -- extra test (freeze)
        expected <- freezeTM tm
        actual <-
          if even n
            then VU.map (\(S (!_, !_, !t)) -> t) <$> LST.freeze seg0
            else VU.map (\(S (!_, !_, !t)) -> t) <$> LST.unsafeFreeze seg0
        liftIO $ actual @?= expected
  
  -- | prod, read, applyIn, applyAt
  unit_naiveTest :: TestTree
  unit_naiveTest = testCase "naiveTest" $
    testDriver (0, 3) 3000 $ \now seg0 tm ty l r -> case ty of
      0 -> do
        -- prod
        S (!resL, !resR, !resTime) <- LST.prod seg0 l r
        liftIO $ l @=? resL
        liftIO $ r @=? resR
        liftIO . (@=? resTime) =<< prod tm l r
        pure now
      1 -> do
        -- read
        S (!resL, !resR, !resTime) <- LST.read seg0 l
        liftIO $ l @=? resL
        liftIO $ l + 1 @=? resR
        liftIO . (@=? resTime) =<< prod tm l (l + 1)
        pure now
      2 -> do
        -- applyIn
        LST.applyIn seg0 l r (T (now + 1))
        action tm l r (now + 1)
        pure $ now + 1
      3 -> do
        -- applyAt
        LST.applyAt seg0 l (T (now + 1))
        action tm l (l + 1) (now + 1)
        pure $ now + 1
      _ -> error "unreachable"
  
  unit_maxRightTest :: TestTree
  unit_maxRightTest = testCase "maxRightTest" $ do
    testDriver (0, 2) 1000 $ \now seg0 tm ty l r -> case ty of
      0 -> do
        -- maxRight
        !_ <- LST.maxRight seg0 l $ \(S (!lS, !rS, !_)) -> case lS of
          _
            | lS == -1 -> True
            | lS /= l -> error "unreachable"
            -- FIXME: add `maxRightM`
            -- \| sTime /= prod tm l sR
            | otherwise -> rS <= r
        pure now
      _ -> do
        LST.applyIn seg0 l r $ T (now + 1)
        action tm l r $ now + 1
        pure $ now + 1
  
  unit_minLeftTest :: TestTree
  unit_minLeftTest = testCase "minLeftTest" $ do
    testDriver (0, 2) 1000 $ \now seg0 tm ty l r -> case ty of
      0 -> do
        -- minLeft
        !_ <- LST.minLeft seg0 r $ \(S (!lS, !rS, !_)) -> case lS of
          _
            | lS == -1 -> True
            | rS /= r -> error "unreachable"
            -- FIXME: add `maxRightM`
            -- \| sTime /= prod tm sL r
            | otherwise -> l <= lS
        pure now
      _ -> do
        LST.applyIn seg0 l r $ T (now + 1)
        action tm l r $ now + 1
        pure $ now + 1
  
  tests :: [TestTree]
  tests =
    [ unit_naiveTest,
      unit_maxRightTest,
      unit_minLeftTest
    ]
  