{-# LANGUAGE DerivingStrategies #-}

module Tests.Extra.Pdsu (tests) where

import AtCoder.Extra.Pdsu qualified as Pdsu
import Control.Monad (forM)
import Control.Monad.ST (runST)
import Data.Maybe (fromJust)
import Test.Tasty
import Test.Tasty.QuickCheck as QC

prop_merge :: QC.Positive Int -> QC.Gen QC.Property
prop_merge (QC.Positive n) = do
  m <- QC.chooseInt (1, 2 * n)
  es <- QC.vectorOf m $ do
    u <- QC.chooseInt (0, n - 1)
    v <- QC.chooseInt (0, n - 1)
    pure (u, v)
  pure . QC.conjoin $ runST $ do
    dsu <- Pdsu.new @_ @() n id
    forM es $ \(!u, !v) -> do
      r' <- Pdsu.merge dsu u v ()
      r1 <- Pdsu.leader dsu u
      r2 <- Pdsu.leader dsu v
      pure (r1 == r' && r2 == r')

prop_mergeMaybe :: QC.Positive Int -> QC.Gen QC.Property
prop_mergeMaybe (QC.Positive n) = do
  m <- QC.chooseInt (1, 2 * n)
  es <- QC.vectorOf m $ do
    u <- QC.chooseInt (0, n - 1)
    v <- QC.chooseInt (0, n - 1)
    pure (u, v)
  pure . QC.conjoin $ runST $ do
    dsu <- Pdsu.new @_ @() n id
    forM es $ \(!u, !v) -> do
      r1 <- Pdsu.leader dsu u
      r2 <- Pdsu.leader dsu v
      r' <- Pdsu.mergeMaybe dsu u v ()
      if r1 == r2
        then pure $ r' == Nothing
        else do
          r1' <- Pdsu.leader dsu u
          r2' <- Pdsu.leader dsu v
          pure (r1' == fromJust r' && r2' == fromJust r')

tests :: [TestTree]
tests =
  [ QC.testProperty "merge" prop_merge,
    QC.testProperty "mergeMaybe" prop_mergeMaybe
  ]