content: {-# LANGUAGE RecordWildCards #-}
  
  module Tests.Extra.WaveletMatrix.Raw (tests) where
  
  import AtCoder.Extra.Bisect (lowerBound)
  import AtCoder.Extra.WaveletMatrix.Raw qualified as WM
  import Control.Exception (evaluate)
  import Data.IntMap qualified as IM
  import Data.Ord (comparing)
  import Data.Vector.Algorithms.Intro qualified as VAI
  import Data.Vector.Unboxed qualified as VU
  import System.IO.Unsafe (unsafePerformIO)
  import Test.Hspec
  import Test.Tasty
  import Test.Tasty.HUnit
  import Test.Tasty.Hspec
  import Test.Tasty.QuickCheck as QC
  import Tests.Util (intervalGen)
  
  compress :: VU.Vector Int -> VU.Vector Int
  compress xs =
    let dict = VU.uniq $ VU.modify VAI.sort xs
     in VU.map (lowerBound dict) xs
  
  data Init = Init
    { capacity :: {-# UNPACK #-} !Int,
      xs :: !(VU.Vector Int),
      wm :: !WM.RawWaveletMatrix
    }
    deriving (Eq, Show)
  
  instance QC.Arbitrary Init where
    arbitrary = do
      QC.NonNegative n <- QC.arbitrary
      xs <- compress . VU.fromListN n <$> QC.vectorOf n (QC.arbitrary @Int)
      pure $ Init n xs (WM.build n xs)
  
  data Query
    = Access !Int
    | KthSmallestIn !(Int, Int) !Int
    | IKthSmallestIn !(Int, Int) !Int
    | KthLargestIn !(Int, Int) !Int
    | IKthLargestIn !(Int, Int) !Int
    | UnsafeKthSmallestIn !(Int, Int) !Int
    | UnsafeIKthSmallestIn !(Int, Int) !Int
    | UnsafeKthLargestIn !(Int, Int) !Int
    | UnsafeIKthLargestIn !(Int, Int) !Int
    | RankLT !(Int, Int) !Int
    | RankBetween !(Int, Int) !(Int, Int)
    | Rank !(Int, Int) !Int
    | Select !Int
    | SelectKth !Int !Int
    | SelectKthIn !(Int, Int) !Int !Int
    | LookupLE !(Int, Int) !Int
    | LookupLT !(Int, Int) !Int
    | LookupGE !(Int, Int) !Int
    | LookupGT !(Int, Int) !Int
    | AssocsIn !(Int, Int)
    | DescAssocsIn !(Int, Int)
    deriving (Show)
  
  genQuery :: Int -> QC.Gen Query
  genQuery n = do
    QC.oneof
      [ Access <$> QC.chooseInt (-1, n),
        KthSmallestIn <$> lr <*> exc,
        IKthSmallestIn <$> lr <*> exc,
        KthLargestIn <$> lr <*> exc,
        IKthLargestIn <$> lr <*> exc,
        UnsafeKthSmallestIn <$> lr <*> inc,
        UnsafeIKthSmallestIn <$> lr <*> inc,
        UnsafeKthLargestIn <$> lr <*> inc,
        UnsafeIKthLargestIn <$> lr <*> inc,
        RankLT <$> lr <*> val,
        RankBetween <$> lr <*> lr,
        Rank <$> lr <*> val,
        Select <$> val,
        SelectKth <$> exc <*> val,
        SelectKthIn <$> lr <*> exc <*> val,
        LookupLE <$> lr <*> QC.chooseInt (-1, n),
        LookupLT <$> lr <*> QC.chooseInt (-1, n),
        LookupGE <$> lr <*> QC.chooseInt (-1, n),
        LookupGT <$> lr <*> QC.chooseInt (-1, n),
        AssocsIn <$> lr,
        DescAssocsIn <$> lr
      ]
    where
      inc = QC.chooseInt (0, n - 1)
      exc = QC.chooseInt (0, n)
      lr = intervalGen n
      val = QC.chooseInt (-n, n)
  
  -- | Arbitrary return type for the `Query` result.
  data Result
    = I {-# UNPACK #-} !Int
    | M !(Maybe Int)
    | M2 !(Maybe (Int, Int))
    | Assocs [(Int, Int)]
    deriving (Show, Eq)
  
  -- | containers. (referential implementation)
  handleRef :: VU.Vector Int -> Query -> Result
  handleRef xs q = case q of
    Access i -> M $ xs VU.!? i
    KthSmallestIn (!l, !r) k -> M $ snd <$> ikthSmallestIn l r k
    IKthSmallestIn (!l, !r) k -> M2 $ ikthSmallestIn l r k
    KthLargestIn (!l, !r) k -> M $ snd <$> ikthLargestIn l r k
    IKthLargestIn (!l, !r) k -> M2 $ ikthLargestIn l r k
    UnsafeKthSmallestIn (!l, !r) k -> M $ snd <$> ikthSmallestIn l r k
    UnsafeIKthSmallestIn (!l, !r) k -> M2 $ ikthSmallestIn l r k
    UnsafeKthLargestIn (!l, !r) k -> M $ snd <$> ikthLargestIn l r k
    UnsafeIKthLargestIn (!l, !r) k -> M2 $ ikthLargestIn l r k
    RankLT (!l, !r) xr -> rankBetween l r (minBound `div` 2) xr
    RankBetween (!l, !r) (!xl, !xr) -> rankBetween l r xl xr
    Rank (!l, !r) x -> rankBetween l r x (x + 1)
    Select x -> M $ selectKthIn 0 n 0 x
    SelectKth k x -> M $ selectKthIn 0 n k x
    SelectKthIn (!l, !r) k x -> M $ selectKthIn l r k x
    LookupLE (!l, !r) x -> max_ . VU.filter (<= x) . VU.take (r - l) $ VU.drop l xs
    LookupLT (!l, !r) x -> max_ . VU.filter (< x) . VU.take (r - l) $ VU.drop l xs
    LookupGE (!l, !r) x -> min_ . VU.filter (>= x) . VU.take (r - l) $ VU.drop l xs
    LookupGT (!l, !r) x -> min_ . VU.filter (> x) . VU.take (r - l) $ VU.drop l xs
    AssocsIn (!l, !r) -> Assocs . IM.assocs . IM.fromListWith (+) . VU.toList . VU.map (,1) . VU.take (r - l) $ VU.drop l xs
    DescAssocsIn (!l, !r) -> Assocs . reverse . IM.assocs . IM.fromListWith (+) . VU.toList . VU.map (,1) . VU.take (r - l) $ VU.drop l xs
    where
      n = VU.length xs
      max_ ys
        | VU.null ys = M Nothing
        | otherwise = M $ Just $ VU.maximum ys
      min_ ys
        | VU.null ys = M Nothing
        | otherwise = M $ Just $ VU.minimum ys
      ikthSmallestIn l r k =
        (VU.!? k)
          . VU.modify (VAI.sortBy (comparing (\(!i, !x) -> (x, i))))
          . VU.take (r - l)
          . VU.drop l
          . VU.indexed
          $ xs
      ikthLargestIn l r k = ikthSmallestIn l r ((r - l) - (k + 1))
      rankBetween l r xl xr =
        I
          . VU.length
          . VU.filter (\x -> xl <= x && x < xr)
          . VU.take (r - l)
          . VU.drop l
          $ xs
      selectKthIn l r k x =
        (fst <$>)
          . (VU.!? k)
          . VU.filter ((== x) . snd)
          . VU.take (r - l)
          . VU.drop l
          . VU.indexed
          $ xs
  
  handleAcl :: WM.RawWaveletMatrix -> Query -> Result
  handleAcl wm q = case q of
    Access i -> M $ WM.access wm i
    KthSmallestIn (!l, !r) k -> M $ WM.kthSmallestIn wm l r k
    IKthSmallestIn (!l, !r) k -> M2 $ WM.ikthSmallestIn wm l r k
    KthLargestIn (!l, !r) k -> M $ WM.kthLargestIn wm l r k
    IKthLargestIn (!l, !r) k -> M2 $ WM.ikthLargestIn wm l r k
    UnsafeKthSmallestIn (!l, !r) k -> M $ WM.kthSmallestIn wm l r k
    UnsafeIKthSmallestIn (!l, !r) k -> M2 $ WM.ikthSmallestIn wm l r k
    UnsafeKthLargestIn (!l, !r) k -> M $ WM.kthLargestIn wm l r k
    UnsafeIKthLargestIn (!l, !r) k -> M2 $ WM.ikthLargestIn wm l r k
    RankLT (!l, !r) xr -> I $ WM.rankLT wm l r xr
    RankBetween (!l, !r) (!xl, !xr) -> I $ WM.rankBetween wm l r xl xr
    Rank (!l, !r) x -> I $ WM.rank wm l r x
    Select x -> M $ WM.select wm x
    SelectKth k x -> M $ WM.selectKth wm k x
    SelectKthIn (!l, !r) k x -> M $ WM.selectKthIn wm l r k x
    LookupLE (!l, !r) x -> M $ WM.lookupLE wm l r x
    LookupLT (!l, !r) x -> M $ WM.lookupLT wm l r x
    LookupGE (!l, !r) x -> M $ WM.lookupGE wm l r x
    LookupGT (!l, !r) x -> M $ WM.lookupGT wm l r x
    AssocsIn (!l, !r) -> Assocs $ WM.assocsIn wm l r
    DescAssocsIn (!l, !r) -> Assocs $ WM.descAssocsIn wm l r
  
  prop_randomTest :: Init -> QC.Gen QC.Property
  prop_randomTest Init {..} = do
    qs <- QC.vectorOf capacity (genQuery capacity)
    pure . QC.conjoin $
      map
        ( \q ->
            QC.counterexample (show q) $
              handleRef xs q QC.=== handleAcl wm q
        )
        qs
  
  unit_boundary :: TestTree
  unit_boundary = testCase "boundary" $ do
    let n = 5
    let wm = WM.build 3 $ VU.fromList [0, 1, 2, 1, 0]
  
    let try :: (HasCallStack, Eq a, Show a) => (WM.RawWaveletMatrix -> Int -> Int -> Int -> Maybe a) -> Int -> IO ()
        try f x = do
          (@?= Nothing) $ f wm (-1) 0 x
          (@?= Nothing) $ f wm n (n + 1) x
  
    let k = 0
    try WM.kthSmallestIn k
    try WM.ikthSmallestIn k
    try WM.kthLargestIn k
    try WM.ikthLargestIn k
  
    -- try WM.unsafeKthSmallestIn
    -- try WM.unsafeIKthSmallestIn
    -- try WM.unsafeKthLargestIn
    -- try WM.unsafeIKthLargestIn
  
    let tryRank :: (HasCallStack) => (WM.RawWaveletMatrix -> Int -> Int -> Int -> Int) -> Int -> IO ()
        tryRank f x = do
          -- out of range
          (@?= 0) $ f wm (-1) 0 x
          (@?= 0) $ f wm n (n + 1) x
          -- reverse
          (@?= 0) $ f wm 1 0 x
          (@?= 0) $ f wm n (n - 1) x
          (@?= 0) $ f wm n 0 x
          -- out of range and reverse
          (@?= 0) $ f wm 0 (-1) x
          (@?= 0) $ f wm (n + 1) n x
  
    let tryRankBetween :: (HasCallStack) => (WM.RawWaveletMatrix -> Int -> Int -> Int -> Int -> Int) -> Int -> Int -> IO ()
        tryRankBetween f xl xr = do
          -- out of range
          (@?= 0) $ f wm (-1) 0 xl xr
          (@?= 0) $ f wm n (n + 1) xl xr
          -- reverse
          (@?= 0) $ f wm 1 0 xl xr
          (@?= 0) $ f wm n (n - 1) xl xr
          (@?= 0) $ f wm n 0 xl xr
          -- out of range and reverse
          (@?= 0) $ f wm 0 (-1) xl xr
          (@?= 0) $ f wm (n + 1) n xl xr
  
    tryRank WM.rankLT 0
    tryRank WM.rankLT 1
    tryRank WM.rankLT (-1)
    tryRankBetween WM.rankBetween (-1) 1
    tryRank WM.rank 0
    tryRank WM.rank 1
    tryRank WM.rank (-1)
  
  -- TODO: test
  
  -- (@?= Nothing) $ WM.select wm
  -- (@?= Nothing) $ WM.selectKth wm
  -- (@?= Nothing) $ WM.selectKthIn wm
  
  -- (@?= Nothing) $ WM.lookupLE wm
  -- (@?= Nothing) $ WM.lookupLT wm
  -- (@?= Nothing) $ WM.lookupGE wm
  -- (@?= Nothing) $ WM.lookupGT wm
  -- (@?= Nothing) $ WM.assocsIn wm
  -- (@?= Nothing) $ WM.descAssocsIn wm
  
  spec_invalid :: IO TestTree
  spec_invalid = testSpec "boundary check" $ do
    it "throws error 1" $ do
      evaluate (WM.build (-1) (VU.singleton 1)) `shouldThrow` anyException
  
  tests :: [TestTree]
  tests =
    [ unit_boundary,
      unsafePerformIO spec_invalid,
      QC.testProperty "random test" prop_randomTest
    ]
  