module Tests.String (tests) where

import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.Internal.String qualified as ACIS
import AtCoder.String qualified as AS
import Control.Monad (when)
import Control.Monad.Fix (fix)
import Control.Monad.ST (runST)
import Data.ByteString.Char8 qualified as BS
import Data.Char (ord)
import Data.Foldable
import Data.Vector qualified as V
import Data.Vector.Algorithms.Intro qualified as VAI
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

saNaive :: VU.Vector Int -> VU.Vector Int
saNaive s = VU.create $ do
  let n = VU.length s
  sa <- VUM.generate n id
  let f l r = compare (VU.drop l s) (VU.drop r s)
  VAI.sortBy f sa
  pure sa

lcpNaive :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
lcpNaive bs sa = VU.create $ do
  let n = VU.length bs
  let !_ = ACIA.runtimeAssert (n > 0) ""
  lcp <- VUM.replicate (n - 1) (0 :: Int)
  for_ [0 .. n - 2] $ \i -> do
    let l = sa VG.! i
    let r = sa VG.! (i + 1)
    fix $ \loop -> do
      lcpI <- VGM.read lcp i
      let sl = bs VU.! (l + lcpI)
      let sr = bs VU.! (r + lcpI)
      when (l + lcpI < n && r + lcpI < n && sl == sr) $ do
        VGM.modify lcp (+ 1) i
        loop
  pure lcp

zNaive :: VU.Vector Int -> VU.Vector Int
zNaive s = VU.create $ do
  let n = VU.length s
  z <- VUM.replicate n (0 :: Int)
  for_ [0 .. n - 1] $ \i -> do
    fix $ \loop -> do
      zi <- VGM.read z i
      let sl = s VU.! zi
      let sr = s VU.! (i + zi)
      when (i + zi < n && sl == sr) $ do
        VGM.modify z (+ 1) i
        loop
  pure z

unit_empty :: TestTree
unit_empty = testCase "empty" $ do
  VU.empty @=? AS.suffixArrayBS BS.empty
  VU.empty @=? AS.suffixArrayOrd (VU.empty @Int)

  VU.empty @=? AS.zAlgorithmBS BS.empty
  VU.empty @=? AS.zAlgorithm (VU.empty @Int)

testWithNaive :: (VU.Vector Int -> VU.Vector Int -> Int -> Assertion) -> Assertion
testWithNaive runTest = do
  for_ [1 .. 5] $ \n -> do
    let m = 4 ^ n :: Int
    for_ [0 .. m - 1] $ \f -> do
      let (!str, !maxC') = runST $ do
            s <- VUM.replicate n (0 :: Int)
            (!_, !maxC) <-
              VU.foldM'
                ( \(!g, !maxC_) i -> do
                    VGM.write s i $ g `mod` 4
                    pure (g `div` 4, max maxC_ (g `mod` 4))
                )
                (f, 0)
                (VU.generate n id)
            (,maxC) <$> VU.unsafeFreeze s
      runTest str (saNaive str) maxC'

  for_ [1 .. 10] $ \n -> do
    let m = 2 ^ n :: Int
    for_ [0 .. m - 1] $ \f -> do
      let (!str, !maxC') = runST $ do
            s <- VUM.replicate n (0 :: Int)
            (!_, !maxC) <-
              VU.foldM'
                ( \(!g, !maxC_) i -> do
                    VGM.write s i $ g `mod` 2
                    pure (g `div` 2, max maxC_ (g `mod` 2))
                )
                (f, 0)
                (VU.generate n id)
            (,maxC) <$> VU.unsafeFreeze s
      runTest str (saNaive str) maxC'

unit_saLcpNaive :: TestTree
unit_saLcpNaive = testCase "saLcpNaive" $ do
  testWithNaive $ \s sa upper -> do
    sa @=? AS.suffixArrayOrd s
    sa @=? AS.suffixArray s upper
    lcpNaive s sa @=? AS.lcpArray s sa

unit_saNaiveNaive :: TestTree
unit_saNaiveNaive = testCase "saNaiveNaive" $ do
  testWithNaive $ \s sa _upper -> do
    sa @=? ACIS.saNaive s

unit_saDoublingNaive :: TestTree
unit_saDoublingNaive = testCase "saDoublingNaive" $ do
  testWithNaive $ \s sa _upper -> do
    sa @=? ACIS.saDoubling s

unit_saIsNaive :: TestTree
unit_saIsNaive = testCase "saIsNaive" $ do
  testWithNaive $ \s sa upper -> do
    sa @=? ACIS.saIsManual (-1) (-1) s upper

unit_allATest :: TestTree
unit_allATest = testCase "allATest" $ do
  for_ [1 .. 100] $ \n -> do
    let s = VU.replicate n (0 :: Int)
    saNaive s @=? AS.suffixArrayOrd s
    saNaive s @=? AS.suffixArray s 10
    saNaive s @=? AS.suffixArray s 12

unit_allBTest :: TestTree
unit_allBTest = testCase "allBTest" $ do
  for_ [1 .. 100] $ \n -> do
    let s = VU.generate n (`mod` 2)
    saNaive s @=? AS.suffixArrayOrd s
    saNaive s @=? AS.suffixArray s 3
  for_ [1 .. 100] $ \n -> do
    let s = VU.generate n ((1 -) . (`mod` 2))
    saNaive s @=? AS.suffixArrayOrd s
    saNaive s @=? AS.suffixArray s 3

-- TODO: all A test
-- TODO: all B test

unit_sa :: TestTree
unit_sa = testCase "sa" $ do
  let s = BS.pack "missisippi"
  let ans =
        V.map BS.pack $
          V.fromList
            [ "i",
              "ippi",
              "isippi",
              "issisippi",
              "missisippi",
              "pi",
              "ppi",
              "sippi",
              "sisippi",
              "ssisippi"
            ]

  testSA "public" s ans $ AS.suffixArrayBS s
  testSA "internal naive" s ans $ ACIS.saNaive (bsToSA s)
  testSA "internal doubling" s ans $ ACIS.saDoubling (bsToSA s)
  testSA "internal SA-IS" s ans $ ACIS.saIsImpl 10 40 (bsToSA s) 255
  where
    testSA msg s expected result = do
      VU.length result @?= V.length expected
      V.iforM_ expected $ \i ansI -> do
        let res = BS.drop (result VG.! i) s
        assertBool (msg ++ ": " ++ show result) (res == ansI)
    bsToSA :: BS.ByteString -> VU.Vector Int
    bsToSA s =
      let n = BS.length s
       in -- FIXME: correct? (upper 255?)
          VU.map ord $ VU.fromListN n (BS.unpack s)

unit_saSingle :: TestTree
unit_saSingle = testCase "saSingle" $ do
  VU.singleton 0 @?= AS.suffixArrayOrd (VU.singleton (0 :: Int))
  VU.singleton 0 @?= AS.suffixArrayOrd (VU.singleton (-1 :: Int))
  VU.singleton 0 @?= AS.suffixArrayOrd (VU.singleton (-1 :: Int))
  VU.singleton 0 @?= AS.suffixArrayOrd (VU.singleton (minBound :: Int))
  VU.singleton 0 @?= AS.suffixArrayOrd (VU.singleton (maxBound :: Int))

unit_lcp :: TestTree
unit_lcp = testCase "LCP" $ do
  let s = BS.pack "aab"
  let sa = AS.suffixArrayBS s
  VU.fromList [0, 1, 2] @=? sa
  let lcp = AS.lcpArrayBS s sa
  VU.fromList [1, 0] @=? lcp

  lcp @=? AS.lcpArray (VU.fromList [0 :: Int, 0, 1]) sa
  lcp @=? AS.lcpArray (VU.fromList [-100 :: Int, -100, 100]) sa
  lcp @=? AS.lcpArray (VU.fromList [minBound :: Int, minBound, maxBound]) sa

unit_zAlgo :: TestTree
unit_zAlgo = testCase "zAlgo" $ do
  let s = BS.pack "abab"
  let z = AS.zAlgorithmBS s
  VU.fromList [4, 0, 2, 0] @=? z
  VU.fromList [4, 0, 2, 0] @=? AS.zAlgorithm (VU.fromList [1 :: Int, 10, 1, 10])
  zNaive (VU.replicate 7 0) @=? AS.zAlgorithm (VU.replicate 7 (0 :: Int))

unit_zNaive :: TestTree
unit_zNaive = testCase "zNaive" $ do
  for_ [1 .. 6] $ \n -> do
    let m = 4 ^ n :: Int
    for_ [0 .. m - 1] $ \f -> do
      let s' = VU.create $ do
            s <- VUM.replicate n (0 :: Int)
            VU.foldM'_
              ( \g i -> do
                  VGM.write s i $ g `mod` 4
                  pure $ g `div` 4
              )
              f
              (VU.generate n id)
            pure s
      zNaive s' @=? AS.zAlgorithm s'

  for_ [1 .. 10] $ \n -> do
    let m = 2 ^ n
    for_ [0 .. m - 1] $ \f -> do
      let s' = VU.create $ do
            s <- VUM.replicate n (0 :: Int)
            VU.foldM'_
              ( \g i -> do
                  VGM.write s i $ g `mod` 2
                  pure $ g `div` 2
              )
              f
              (VU.generate n id)
            pure s
      zNaive s' @=? AS.zAlgorithm s'

spec_invalid :: IO TestTree
spec_invalid = testSpec "invalid" $ do
  -- safe:
  let !_ = AS.suffixArray (VU.fromList [0, 10]) 10
  it "throws error" $ (`shouldThrow` anyException) $ do
    let !_ = AS.suffixArray VU.empty (-1)
    pure ()
  it "throws error" $ (`shouldThrow` anyException) $ do
    let !_ = AS.suffixArray (VU.singleton (-1)) 10
    pure ()
  it "throws error" $ (`shouldThrow` anyException) $ do
    let !_ = AS.suffixArray (VU.singleton 11) 10
    pure ()
  it "throws error" $ (`shouldThrow` anyException) $ do
    let !_ = AS.suffixArray (VU.fromList [0, 11]) 10
    pure ()
  it "throws error" $ (`shouldThrow` anyException) $ do
    let !_ = AS.suffixArray (VU.fromList [0, -1]) 10
    pure ()

  -- safe:
  let !_ = AS.suffixArrayOrd (VU.fromList [-999 :: Int, 999])
  pure ()

tests :: [TestTree]
tests =
  [ unit_empty,
    unit_saLcpNaive,
    unit_saNaiveNaive,
    unit_saDoublingNaive,
    unit_saIsNaive,
    unit_allATest,
    unit_allBTest,
    unit_sa,
    unit_saSingle,
    unit_lcp,
    unit_zAlgo,
    unit_zNaive,
    unsafePerformIO spec_invalid
  ]