content: {-# LANGUAGE DerivingVia #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE TypeFamilies #-}
  
  module Tests.SegTree (tests) where
  
  import AtCoder.Internal.Assert
  import AtCoder.SegTree qualified as ST
  import Control.Monad.Primitive (PrimMonad, PrimState)
  import Data.Char (chr, ord)
  import Data.Foldable (for_)
  import Data.Monoid
  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.IO.Unsafe (unsafePerformIO)
  import Test.Hspec
  import Test.Tasty
  import Test.Tasty.HUnit
  import Test.Tasty.Hspec
  
  data SegTreeNaive s a = SegTreeNaive
    { nStn :: {-# UNPACK #-} !Int,
      dStn :: !(VUM.MVector s a)
    }
  
  newStn :: (PrimMonad m, Monoid a, VU.Unbox a) => Int -> m (SegTreeNaive (PrimState m) a)
  newStn nStn = do
    dStn <- VUM.replicate nStn mempty
    pure SegTreeNaive {..}
  
  writeStn :: (VU.Unbox a, PrimMonad m) => SegTreeNaive (PrimState m) a -> Int -> a -> m ()
  writeStn SegTreeNaive {..} = VGM.write dStn
  
  readStn :: (VU.Unbox a, PrimMonad m) => SegTreeNaive (PrimState m) a -> Int -> m a
  readStn SegTreeNaive {..} = VGM.read dStn
  
  prodStn :: (Monoid a, VU.Unbox a, PrimMonad m) => SegTreeNaive (PrimState m) a -> Int -> Int -> m a
  prodStn SegTreeNaive {..} l r = do
    d <- VU.unsafeFreeze dStn
    let slice = VU.take (r - l) $ VU.drop l d
    pure $ VU.foldl' (<>) mempty slice
  
  maxRightStn :: (HasCallStack, Monoid a, VU.Unbox a, PrimMonad m) => SegTreeNaive (PrimState m) a -> Int -> (a -> Bool) -> m Int
  maxRightStn SegTreeNaive {..} l0 f = do
    let !_ = runtimeAssert (f mempty) "ident"
    let loop _ _ [] = nStn
        loop l !acc (x : xs)
          | not (f acc') = l
          | otherwise = loop (l + 1) acc' xs
          where
            !acc' = acc <> x
    d <- VU.unsafeFreeze dStn
    pure . loop l0 mempty . VU.toList $ VU.drop l0 d
  
  minLeftStn :: (HasCallStack, Monoid a, VU.Unbox a, PrimMonad m) => SegTreeNaive (PrimState m) a -> Int -> (a -> Bool) -> m Int
  minLeftStn SegTreeNaive {..} r0 f = do
    let !_ = runtimeAssert (f mempty) "ident"
    let loop _ _ [] = 0
        loop r !acc (x : xs)
          | not (f acc') = r + 1
          | otherwise = loop (r - 1) acc' xs
          where
            !acc' = x <> acc
    d <- VU.unsafeFreeze dStn
    pure . loop (r0 - 1) mempty . VU.toList . VU.reverse $ VU.take r0 d
  
  freezeStn :: (Monoid a, VU.Unbox a, PrimMonad m) => SegTreeNaive (PrimState m) a -> m (VU.Vector a)
  freezeStn = VU.freeze . dStn
  
  type FooRepr = VU.DoNotUnboxStrict String
  
  newtype Foo = Foo String
    deriving (Eq, Ord, Show)
  
  instance Semigroup Foo where
    (<>) :: (HasCallStack) => Foo -> Foo -> Foo
    (Foo a) <> (Foo b) = Foo $ case (a, b) of
      ("$", _) -> b
      (_, "$") -> a
      _ -> a <> b
      where
        !_ = runtimeAssert (a == "$" || b == "$" || a <= b) "error"
  
  instance Monoid Foo where
    mempty = Foo "$"
  
  instance VU.IsoUnbox Foo FooRepr where
    {-# INLINE toURepr #-}
    toURepr (Foo s) = VU.DoNotUnboxStrict s
    {-# INLINE fromURepr #-}
    fromURepr (VU.DoNotUnboxStrict s) = Foo s
  
  newtype instance VU.MVector s Foo = MV_Foo (VU.MVector s FooRepr)
  
  newtype instance VU.Vector Foo = V_Foo (VU.Vector FooRepr)
  
  deriving via (Foo `VU.As` FooRepr) instance VGM.MVector VUM.MVector Foo
  
  deriving via (Foo `VU.As` FooRepr) instance VG.Vector VU.Vector Foo
  
  instance VU.Unbox Foo
  
  unit_zero :: TestTree
  unit_zero = testCase "zero" $ do
    s <- ST.new @_ @(Sum Int) 0
    (@?= mempty) =<< ST.allProd s
  
    pure ()
  
  spec_invalid :: IO TestTree
  spec_invalid = testSpec "invalid" $ do
    it "throws error" $
      ST.new @_ @(Sum Int) (-1) `shouldThrow` anyException
    s <- runIO $ ST.new @_ @(Sum Int) 10
  
    it "throws error" $ ST.read s (-1) `shouldThrow` anyException
    it "throws error" $ ST.read s 10 `shouldThrow` anyException
  
    it "throws error" $ ST.write s (-1) 0 `shouldThrow` anyException
    it "throws error" $ ST.write s 10 0 `shouldThrow` anyException
    it "throws error" $ ST.modify s (+ 1) (-1) `shouldThrow` anyException
    it "throws error" $ ST.modify s (+ 1) 10 `shouldThrow` anyException
    it "throws error" $ ST.modifyM s (pure . (+ 1)) (-1) `shouldThrow` anyException
    it "throws error" $ ST.modifyM s (pure . (+ 1)) 10 `shouldThrow` anyException
  
    it "throws error" $ ST.prod s (-1) (-1) `shouldThrow` anyException
    it "throws error" $ ST.prod s 3 2 `shouldThrow` anyException
    it "throws error" $ ST.prod s 0 11 `shouldThrow` anyException
    it "throws error" $ ST.prod s (-1) 11 `shouldThrow` anyException
  
    it "throws error" $ ST.maxRight s 11 (const True) `shouldThrow` anyException
    it "throws error" $ ST.minLeft s (-1) (const True) `shouldThrow` anyException
    it "throws error" $ ST.maxRight s 0 (const False) `shouldThrow` anyException
  
  unit_one :: TestTree
  unit_one = testCase "one" $ do
    seg <- ST.new @_ @(Sum Int) 1
    (@?= mempty) =<< ST.allProd seg
    (@?= mempty) =<< ST.read seg 0
    (@?= mempty) =<< ST.prod seg 0 1
    let dummy = Sum 42
    ST.write seg 0 dummy
    (@?= dummy) =<< ST.read seg 0
    (@?= mempty) =<< ST.prod seg 0 0
    (@?= dummy) =<< ST.prod seg 0 1
    (@?= mempty) =<< ST.prod seg 1 1
  
  unit_compareNaive :: TestTree
  unit_compareNaive = testCase "compareNaive" $ do
    run $ \seg i x -> ST.write seg i x
    run $ \seg i x -> ST.modify seg (const x) i
    run $ \seg i x -> ST.modifyM seg (pure . (const x)) i
    where
      run f = do
        for_ [0 .. 30 - 1] $ \n -> do
          seg0 <- newStn @_ @Foo n
          seg1 <- ST.new @_ @Foo n
          -- write
          for_ [0 .. n - 1] $ \i -> do
            writeStn seg0 i . Foo . (: []) . chr $ ord 'a' + i
            (f seg1 i) . Foo . (: []) . chr $ ord 'a' + i
  
          -- prod
          for_ [0 .. n] $ \l -> do
            for_ [l .. n] $ \r -> do
              x0 <- prodStn seg0 l r
              x1 <- ST.prod seg1 l r
              assertEqual (show (l, r)) x0 x1
  
          -- allProd
          do
            x1 <- prodStn seg0 0 n
            x2 <- ST.prod seg1 0 n
            x3 <- ST.allProd seg1
            x1 @?= x2
            x1 @?= x3
  
          -- maxRight
          for_ [0 .. n] $ \l -> do
            for_ [l .. n] $ \r -> do
              Foo y <- ST.prod seg1 l r
              let p :: Foo -> Bool
                  p (Foo x) = length x <= length y
              r0 <- maxRightStn seg0 l p
              r1 <- ST.maxRight seg1 l p
              assertEqual (show ((l, r), y)) r0 r1
  
          -- minLeft
          for_ [0 .. n] $ \r -> do
            for_ [0 .. r] $ \l -> do
              Foo y <- ST.prod seg1 l r
              let p :: Foo -> Bool
                  p (Foo x) = length x <= length y
              r0 <- minLeftStn seg0 r p
              r1 <- ST.minLeft seg1 r p
              assertEqual (show ((l, r), y)) r0 r1
  
          -- read (extra test)
          for_ [0 .. n - 1] $ \i -> do
            expected <- readStn seg0 i
            (@?= expected) =<< ST.read seg1 i
  
          -- freeze (extra test)
          expected <- freezeStn seg0
          (@?= expected) =<< ST.freeze seg1
  
  unit_freezeZero :: TestTree
  unit_freezeZero = testCase "freezeZero" $ do
    seg <- ST.new @_ @(Sum Int) 0
    (VU.empty @=?) =<< ST.freeze seg
    (VU.empty @=?) =<< ST.unsafeFreeze seg
  
  unit_prodMaybeBounds :: TestTree
  unit_prodMaybeBounds = testCase "prodMaybeBounds" $ do
    seg <- ST.new @_ @(Sum Int) 4
    (@?= Just (Sum {getSum = 0})) =<< ST.prodMaybe seg 0 0
    (@?= Just (Sum {getSum = 0})) =<< ST.prodMaybe seg 0 4
    (@?= Just (Sum {getSum = 0})) =<< ST.prodMaybe seg 4 4
    (@?= Nothing) =<< ST.prodMaybe seg (-1) 4
    (@?= Nothing) =<< ST.prodMaybe seg 0 5
    (@?= Nothing) =<< ST.prodMaybe seg 0 (-1)
    (@?= Nothing) =<< ST.prodMaybe seg (-1) (-1)
    (@?= Nothing) =<< ST.prodMaybe seg (-1) 0
    (@?= Nothing) =<< ST.prodMaybe seg 4 5
    (@?= Nothing) =<< ST.prodMaybe seg 5 5
  
  tests :: [TestTree]
  tests =
    [ unit_zero,
      unsafePerformIO spec_invalid,
      unit_one,
      unit_compareNaive,
      unit_freezeZero,
      unit_prodMaybeBounds
    ]
  