content: {-# LANGUAGE ViewPatterns #-}
  
  module Tests.Extra.Vector.Prim where
  
  import AtCoder.Extra.Vector.Prim qualified as EV
  import Control.Monad.ST (ST, runST)
  import Data.Vector.Generic qualified as VG
  import Data.Vector.Unboxed qualified as VU
  import Test.Tasty
  import Test.Tasty.QuickCheck as QC
  
  prop_replicateM :: Int -> Int -> QC.Property
  prop_replicateM n x = VU.replicate n x QC.=== runST (EV.replicateM n (pure x))
  
  prop_generateM :: Int -> Int -> QC.Property
  prop_generateM n x = VU.generate n (+ x) QC.=== runST (EV.generateM n (pure . (+ x)))
  
  prop_iterateNM :: Int -> Int -> QC.Property
  prop_iterateNM n x = VU.iterateN n (* 2) x QC.=== runST (EV.iterateNM n (pure . (* 2)) x)
  
  prop_constructNM :: QC.NonNegative Int -> QC.Property
  prop_constructNM (QC.NonNegative n) = VU.constructN n VG.length QC.=== runST (EV.constructNM n (pure . VG.length))
  
  prop_constructrNM :: QC.NonNegative Int -> QC.Property
  prop_constructrNM (QC.NonNegative n) = VU.constructrN n VG.length QC.=== runST (EV.constructrNM n (pure . VG.length))
  
  prop_mapM :: [Int] -> QC.Property
  prop_mapM (VU.fromList -> xs) = VU.map (* 2) xs QC.=== runST (EV.mapM (pure . (* 2)) xs)
  
  -- prop_mapM_ :: [Int] -> QC.Property
  -- prop_mapM_ xs =
  
  prop_imapM :: [Int] -> QC.Property
  prop_imapM (VU.fromList -> xs) = VU.imap (\i x -> 2 * (i + x)) xs QC.=== runST (EV.imapM (\i x -> pure (2 * (i + x))) xs)
  
  -- prop_imapM_ :: [Int] -> QC.Property
  -- prop_imapM_ xs =
  
  prop_iforM :: [Int] -> QC.Property
  prop_iforM (VU.fromList -> xs) = VU.imap (\i x -> 2 * (i + x)) xs QC.=== runST (EV.iforM xs (\i x -> pure (2 * (i + x))))
  
  -- prop_iforM_ :: [Int] -> QC.Property
  -- prop_iforM_ xs =
  
  prop_zipWithM :: [Int] -> [Int] -> QC.Property
  prop_zipWithM (VU.fromList -> xs) (VU.fromList -> ys) =
    VU.zipWith (*) xs ys QC.=== runST (EV.zipWithM (\x y -> pure (x * y)) xs ys)
  
  -- prop_zipWithM_ :: [Int] -> [Int] -> QC.Property
  -- prop_zipWithM_  (VU.fromList -> xs) (VU.fromList -> ys) =
  
  prop_izipWithM :: [Int] -> [Int] -> QC.Property
  prop_izipWithM (VU.fromList -> xs) (VU.fromList -> ys) =
    VU.izipWith (\i a b -> i + a * b) xs ys QC.=== runST (EV.izipWithM (\i a b -> pure (i + a * b)) xs ys)
  
  -- prop_izipWithM_ :: [Int] -> [Int] -> QC.Property
  -- prop_izipWithM_  (VU.fromList -> xs) (VU.fromList -> ys) =
  
  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_filterM :: [Int] -> QC.Property
  prop_filterM (VU.fromList -> xs) = VU.filter odd xs QC.=== runST (EV.filterM (pure . odd) xs)
  
  prop_mapMaybeM :: [Int] -> QC.Property
  prop_mapMaybeM (VU.fromList -> xs) =
    let f x = if odd x then Just (2 * x) else Nothing
     in VU.mapMaybe f xs QC.=== runST (EV.mapMaybeM (pure . f) xs)
  
  prop_imapMaybeM :: [Int] -> QC.Property
  prop_imapMaybeM (VU.fromList -> xs) =
    let f i x = if odd i then Just (i * x) else Nothing
     in VU.imapMaybe f xs QC.=== runST (EV.imapMaybeM (\i -> pure . f i) xs)
  
  -- | scanM etc.
  prop_monadicScanlLike ::
    ((Int -> Int -> Int) -> Int -> VU.Vector Int -> VU.Vector Int) ->
    (forall s. (Int -> Int -> ST s Int) -> Int -> VU.Vector Int -> ST s (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.=== runST (acl mf x xs')
  
  -- | scanM1 etc.
  prop_monadicScanl1Like ::
    ((Int -> Int -> Int) -> VU.Vector Int -> VU.Vector Int) ->
    (forall s. (Int -> Int -> ST s Int) -> VU.Vector Int -> ST s (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.=== runST (acl mf xs')
  
  tests :: [TestTree]
  tests =
    [ QC.testProperty "replicateM" prop_replicateM,
      QC.testProperty "generateM" prop_generateM,
      QC.testProperty "iterateNM" prop_iterateNM,
      QC.testProperty "constructNM" prop_constructNM,
      QC.testProperty "constructrNM" prop_constructrNM,
      QC.testProperty "mapM" prop_mapM,
      -- QC.testProperty "mapM_" prop_mapM_,
      QC.testProperty "imapM" prop_imapM,
      -- QC.testProperty "imapM_" prop_imapM_,
      QC.testProperty "iforM" prop_iforM,
      -- QC.testProperty "iforM_" prop_iforM_,
      QC.testProperty "zipWithM" prop_zipWithM,
      -- QC.testProperty "zipWithM_" prop_zipWithM_,
      QC.testProperty "izipWithM" prop_izipWithM,
      -- QC.testProperty "izipWithM_" prop_izipWithM_,
      QC.testProperty "concatMapM" prop_concatMapM,
      QC.testProperty "iconcatMap" prop_iconcatMap,
      QC.testProperty "iconcatMapM" prop_iconcatMapM,
      QC.testProperty "filterM" prop_filterM,
      QC.testProperty "mapMaybeM" prop_mapMaybeM,
      QC.testProperty "imapMaybeM" prop_imapMaybeM,
      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')
    ]
  