content: {-# LANGUAGE RecordWildCards #-}
  
  module Tests.Extra.WaveletMatrix2d (tests) where
  
  import AtCoder.Extra.WaveletMatrix2d qualified as WM
  import Control.Monad (foldM_)
  import Control.Monad.Primitive (PrimMonad, PrimState)
  import Control.Monad.ST (RealWorld)
  import Data.List qualified as L
  import Data.Map qualified as M
  import Data.Maybe (fromMaybe)
  import Data.Semigroup (Sum (..))
  import Data.Vector.Unboxed qualified as VU
  import Test.QuickCheck.Monadic as QCM
  import Test.Tasty
  import Test.Tasty.QuickCheck as QC
  import Tests.Util (intervalGen)
  
  data Init = Init
    { capacity :: {-# UNPACK #-} !Int,
      map0 :: !(M.Map (Int, Int) (Sum Int)),
      wmM :: !(IO (WM.WaveletMatrix2d RealWorld (Sum Int)))
    }
  
  instance Show Init where
    show Init {..} = show ("Init", capacity, map0)
  
  instance QC.Arbitrary Init where
    arbitrary = do
      QC.NonNegative n <- QC.arbitrary
      let yxs = VU.fromList [(x, y) | x <- [-16 .. 16], y <- [-16 .. 16]]
      pure $ Init n M.empty (WM.new negate yxs)
  
  data Query
    = Read !(Int, Int)
    | Write !(Int, Int) !Int
    | ModifyAdd !Int !(Int, Int)
    | Prod !(Int, Int) !(Int, Int)
    | ProdMaybe !(Int, Int) !(Int, Int)
    | AllProd
    deriving (Show)
  
  genQuery :: Int -> QC.Gen Query
  genQuery n = do
    QC.oneof
      [ Read <$> lr,
        Write <$> lr <*> val,
        ModifyAdd <$> val <*> lr,
        Prod <$> lr <*> lr,
        ProdMaybe <$> lr <*> lr,
        pure AllProd
      ]
    where
      lr = (\(!x, !y) -> (x - 16, y - 16)) <$> intervalGen 32
      val = QC.arbitrary @Int
  
  -- | Arbitrary return type for the `Query` result.
  data Result
    = None
    | S !(Sum Int)
    | MS !(Maybe (Sum Int))
    deriving (Show, Eq)
  
  -- | containers. (referential implementation)
  handleRef :: M.Map (Int, Int) (Sum Int) -> Query -> (Result, M.Map (Int, Int) (Sum Int))
  handleRef map q = case q of
    Read (!x, !y) -> (S . fromMaybe mempty $ M.lookup (x, y) map, map)
    Write (!x, !y) v -> (None, M.insert (x, y) (Sum v) map)
    ModifyAdd w (!x, !y) -> (None, M.insertWith (+) (x, y) (Sum w) map)
    Prod (!x1, !x2) (!y1, !y2) -> (S $ prod x1 x2 y1 y2, map)
    ProdMaybe (!x1, !x2) (!y1, !y2) -> (MS . Just $ prod x1 x2 y1 y2, map)
    AllProd -> (S $ L.foldl' (<>) mempty (M.elems map), map)
    where
      prod x1 x2 y1 y2 =
        L.foldl' (<>) (mempty :: Sum Int)
          . (snd <$>)
          . filter (\((!x, !y), !_) -> x1 <= x && x < x2 && y1 <= y && y < y2)
          $ M.assocs map
  
  handleAcl :: (PrimMonad m) => WM.WaveletMatrix2d (PrimState m) (Sum Int) -> Query -> m Result
  handleAcl wm q = case q of
    Read (!x, !y) -> do
      S <$> WM.read wm (x, y)
    Write (!x, !y) v -> do
      WM.write wm (x, y) $ Sum v
      pure None
    ModifyAdd w (!x, !y) -> do
      WM.modify wm (+ Sum w) (x, y)
      pure None
    Prod (!x1, !x2) (!y1, !y2) -> do
      S <$> WM.prod wm x1 x2 y1 y2
    ProdMaybe (!x1, !x2) (!y1, !y2) -> do
      MS <$> WM.prodMaybe wm x1 x2 y1 y2
    AllProd -> do
      S <$> WM.allProd wm
  
  prop_randomTest :: Init -> QC.Property
  prop_randomTest Init {..} = QCM.monadicIO $ do
    wm <- QCM.run wmM
    qs <- QCM.pick $ QC.vectorOf capacity (genQuery capacity)
    foldM_
      ( \map query -> do
          let (!expected, !map') = handleRef map query
          actual <- QCM.run $ handleAcl wm query
          QCM.assertWith (expected == actual) $ show (query, expected, actual)
          pure map'
      )
      map0
      qs
  
  tests :: [TestTree]
  tests =
    [ -- unit_boundary,
      QC.testProperty "random test" prop_randomTest
    ]
  