content: module Tests.Extra.AhoCorasick where
  
  import AtCoder.Extra.AhoCorasick qualified as Ac
  import Data.Char (ord)
  import Data.List qualified as L
  import Data.Vector qualified as V
  import Data.Vector.Algorithms qualified as VA
  import Data.Vector.Algorithms.Intro qualified as VAI
  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
  
  rng :: (Int, Int)
  rng = (1, 32)
  
  -- | Generates non-empty string as a vector.
  genString :: Gen (VU.Vector Int)
  genString = do
    VU.fromList <$> QC.listOf1 QC.arbitrary
  
  -- | Generates non-empty strings as a vectors.
  -- not empty
  genStrings :: Gen (V.Vector (VU.Vector Int))
  genStrings = do
    n <- QC.chooseInt rng
    V.fromList <$> QC.vectorOf n genString
  
  prop_empty :: QC.Positive Int -> Bool
  prop_empty (QC.Positive n) = do
    -- Not dies
    let !ac0 = Ac.build V.empty
     in True
  
  prop_numberOfNodes :: TestTree
  prop_numberOfNodes = QC.testProperty "numberOfNodes" $ do
    patterns <- genStrings
    let !ac = Ac.build patterns
    pure $ Ac.size ac <= (1 + VG.sum (VG.map VG.length patterns))
  
  unit_banana :: TestTree
  unit_banana = testCase "banana" $ do
    let patterns =
          V.fromList $
            map
              (VU.map ord . VU.fromList)
              ["banana", "benefit", "banapple", "ban", "fit"]
    let text = VU.map ord $ VU.fromList "banana bans, apple benefits."
    --                        positions:
    --                                   0123456789012345678901234567
    --                                   0         1         2
    --                        5 matches:
    --                                   ban*   ban*        benefit*
    --                                   banana*                fit*
  
    let ac = Ac.build patterns
    let matches = Ac.match ac text
    matches @?= VU.fromList [(3, 3), (6, 0), (10, 3), (26, 1), (26, 4)]
  
  genAlphabet :: Gen [Int]
  genAlphabet = do
    n <- QC.chooseInt (1, 26)
    L.nub <$> QC.vectorOf n QC.arbitrary
  
  genPatternIn :: [Int] -> Gen (VU.Vector Int)
  genPatternIn alphabet = do
    n <- QC.chooseInt (1, 6)
    VU.fromList <$> QC.vectorOf n (QC.elements alphabet)
  
  genTextIn :: [Int] -> Gen (VU.Vector Int)
  genTextIn alphabet = do
    n <- QC.chooseInt (0, 1000)
    VU.fromList <$> QC.vectorOf n (QC.elements alphabet)
  
  -- | Returns a vector of @(endPos, patternId)@
  naiveMatch :: (HasCallStack) => V.Vector (VU.Vector Int) -> VU.Vector Int -> VU.Vector (Int, Int)
  naiveMatch patterns text = VU.concatMap f $ VU.generate (VU.length text) (+ 1)
    where
      nPatterns = V.length patterns
      f :: (HasCallStack) => Int -> VU.Vector (Int, Int)
      f endPos = VU.map (endPos,) . VU.filter (test endPos) $ VU.generate nPatterns id
      test endPos iPattern = t == pat
        where
          pat = patterns VG.! iPattern
          len = VU.length pat
          t = VU.drop (endPos - len) $ VU.take endPos text
  
  prop_random :: TestTree
  prop_random = QC.testProperty "random" $ do
    alphabet <- genAlphabet
    n <- QC.chooseInt (1, 100)
    -- NOTE: Patterns must be deduplicated in order to get exact result
    patterns <- VA.nub . V.fromList <$> QC.vectorOf n (genPatternIn alphabet)
    let !ac = Ac.build patterns
    text <- genTextIn alphabet
    pure
      . QC.counterexample (show (patterns, "in", text))
      $ VG.modify VAI.sort (Ac.match ac text) QC.=== VG.modify VAI.sort (naiveMatch patterns text)
  
  tests :: [TestTree]
  tests =
    [ QC.testProperty "empty" prop_empty,
      prop_numberOfNodes,
      unit_banana,
      prop_random
    ]
  