content: module Tests.Extra.Vector where
  
  import AtCoder.Extra.Vector qualified as EV
  import Control.Monad.ST (runST)
  import Data.Functor.Identity (Identity, runIdentity)
  import Data.List qualified as L
  import Data.Vector qualified as V
  import Data.Vector.Generic qualified as VG
  import Data.Vector.Unboxed qualified as VU
  import Test.Tasty
  import Test.Tasty.HUnit
  import Test.Tasty.QuickCheck as QC
  
  prop_argsort :: [Int] -> QC.Property
  prop_argsort xs =
    let lhs = VU.fromList . map snd . L.sort $ zip xs [0 :: Int ..]
        rhs = EV.argsort $ VU.fromList xs
     in lhs QC.=== rhs
  
  prop_concatMapM :: [Int] -> QC.Property
  prop_concatMapM xs =
    let f x = VU.fromList [x, x, x]
        vec = VU.fromList xs
        lhs = VU.concatMap f vec
        rhs = runST $ EV.concatMapM (pure . f) vec
     in lhs QC.=== rhs
  
  prop_iconcatMap :: [Int] -> QC.Property
  prop_iconcatMap xs =
    let f i x = VU.fromList [i + x, i + x, i + x]
        vec = VU.fromList xs
        lhs = VU.concat $ zipWith f [0 :: Int ..] xs
        rhs = runST $ EV.iconcatMapM (\i x -> pure (f i x)) vec
     in lhs QC.=== rhs
  
  prop_iconcatMapM :: [Int] -> QC.Property
  prop_iconcatMapM xs =
    let f i x = VU.fromList [i + x, i + x, i + x]
        vec = VU.fromList xs
        lhs = VU.concat $ zipWith f [0 :: Int ..] xs
        rhs = runST $ EV.iconcatMapM (\i x -> pure (f i x)) vec
     in lhs QC.=== rhs
  
  prop_mapAccumL :: [Int] -> QC.Property
  prop_mapAccumL xs =
    let f s x = (s * x, s + x)
        (!l1, !l2) = L.mapAccumL f (0 :: Int) xs
        (!r1, !r2) = EV.mapAccumL f (0 :: Int) $ VU.fromList xs
     in QC.conjoin [l1 QC.=== r1, VU.fromList l2 QC.=== r2]
  
  -- | scanM etc.
  prop_monadicScanlLike ::
    ((Int -> Int -> Int) -> Int -> VU.Vector Int -> VU.Vector Int) ->
    ((Int -> Int -> Identity Int) -> Int -> VU.Vector Int -> Identity (VU.Vector Int)) ->
    Int ->
    [Int] ->
    QC.Property
  prop_monadicScanlLike ref acl x xs =
    let xs' = VU.fromList xs
        f = (+)
        mf x y = pure $ x + y
     in ref f x xs' QC.=== runIdentity (acl mf x xs')
  
  -- | scanM1 etc.
  prop_monadicScanl1Like ::
    ((Int -> Int -> Int) -> VU.Vector Int -> VU.Vector Int) ->
    ((Int -> Int -> Identity Int) -> VU.Vector Int -> Identity (VU.Vector Int)) ->
    QC.NonEmptyList Int ->
    QC.Property
  prop_monadicScanl1Like ref acl (QC.NonEmpty xs) =
    let xs' = VU.fromList xs
        f = (+)
        mf x y = pure $ x + y
     in ref f xs' QC.=== runIdentity (acl mf xs')
  
  prop_chunks :: QC.Positive Int -> [Int] -> QC.Property
  prop_chunks (QC.Positive k) [] = EV.chunks k (VU.empty @Int) QC.=== V.empty
  prop_chunks (QC.Positive k) xs =
    let res = EV.chunks k $ VU.fromList xs
        n = length xs
     in QC.conjoin
          [ V.sum (VG.map VG.length res) QC.=== n,
            V.all ((== k) . VG.length) (V.init res) QC.=== True,
            VG.concat (V.toList res) QC.=== VU.fromList xs
          ]
  
  unit_maxRangeSum :: TestTree
  unit_maxRangeSum = testCase "unit_maxRangeSum" $ do
    EV.maxRangeSum (VU.singleton (-1 :: Int)) @?= 0
    EV.maxRangeSum (VU.empty @Int) @?= 0
  
  prop_maxRangeSum :: [Int] -> QC.Property
  prop_maxRangeSum xs =
    let vec = VU.fromList xs
        lhs =
          let n = VU.length vec
              lrs = [(l, r) | l <- [0 .. n], r <- [l .. n]]
              eval (!l, !r) = VU.sum . VU.take (r - l) $ VU.drop l vec
           in maximum $ map eval lrs
        rhs = EV.maxRangeSum vec
     in lhs QC.=== rhs
  
  unit_minRangeSum :: TestTree
  unit_minRangeSum = testCase "unit_minRangeSum [1]" $ do
    EV.minRangeSum (VU.singleton (1 :: Int)) @?= 0
    EV.minRangeSum (VU.empty @Int) @?= 0
  
  prop_minRangeSum :: [Int] -> QC.Property
  prop_minRangeSum xs =
    let vec = VU.fromList xs
        lhs =
          let n = VU.length vec
              lrs = [(l, r) | l <- [0 .. n], r <- [l .. n]]
              eval (!l, !r) = VU.sum . VU.take (r - l) $ VU.drop l vec
           in minimum $ map eval lrs
        rhs = EV.minRangeSum vec
     in lhs QC.=== rhs
  
  prop_slideMinIndices :: QC.Positive Int -> [Int] -> QC.Property
  prop_slideMinIndices (QC.Positive k) xs =
    let vec = VU.fromList xs
     in slideMin k vec QC.=== EV.slideMinIndices k vec
  
  prop_slideMaxIndices :: QC.Positive Int -> [Int] -> QC.Property
  prop_slideMaxIndices (QC.Positive k) xs =
    let vec = VU.fromList xs
     in slideMax k vec QC.=== EV.slideMaxIndices k vec
  
  slideMin :: Int -> VU.Vector Int -> VU.Vector Int
  slideMin k xs
    | VU.null xs = VU.empty
    | k >= n = VU.singleton $ VU.minIndex xs
    | otherwise = VU.generate (n - (k - 1)) $ \l ->
        let slice = VU.take k $ VU.drop l xs
         in (+ l) $ VU.minIndex slice
    where
      n = VU.length xs
  
  slideMax :: Int -> VU.Vector Int -> VU.Vector Int
  slideMax k xs
    | VU.null xs = VU.empty
    | k >= n = VU.singleton $ VU.maxIndex xs
    | otherwise = VU.generate (n - (k - 1)) $ \l ->
        let slice = VU.take k $ VU.drop l xs
         in (+ l) $ VU.maxIndex slice
    where
      n = VU.length xs
  
  tests :: [TestTree]
  tests =
    [ QC.testProperty "argsort" prop_argsort,
      QC.testProperty "concatMapM" prop_concatMapM,
      QC.testProperty "iconcatMap" prop_iconcatMap,
      QC.testProperty "iconcatMapM" prop_iconcatMapM,
      QC.testProperty "mapAccumL" prop_mapAccumL,
      QC.testProperty "chunks" prop_chunks,
      QC.testProperty "prescanlM" (prop_monadicScanlLike VU.prescanl EV.prescanlM),
      QC.testProperty "prescanlM'" (prop_monadicScanlLike VU.prescanl' EV.prescanlM'),
      QC.testProperty "postscanlM" (prop_monadicScanlLike VU.postscanl EV.postscanlM),
      QC.testProperty "postscanlM'" (prop_monadicScanlLike VU.postscanl' EV.postscanlM'),
      QC.testProperty "scanlM" (prop_monadicScanlLike VU.scanl EV.scanlM),
      QC.testProperty "scanlM'" (prop_monadicScanlLike VU.scanl' EV.scanlM'),
      QC.testProperty "scanl1M" (prop_monadicScanl1Like VU.scanl1 EV.scanl1M),
      QC.testProperty "scanl1M'" (prop_monadicScanl1Like VU.scanl1' EV.scanl1M'),
      QC.testProperty "maxRangeSum" prop_maxRangeSum,
      QC.testProperty "minRangeSum" prop_minRangeSum,
      unit_maxRangeSum,
      unit_minRangeSum,
      QC.testProperty "slideMinIndices" prop_slideMinIndices,
      QC.testProperty "slideMaxIndices" prop_slideMaxIndices
    ]
  