content: module Tests.Util
    ( myForAllShrink,
      laws,
      intervalGen,
      intervalGen',
      Dag (..),
      genDag,
      Tree (..),
      genTree,
    )
  where
  
  import AtCoder.Extra.Graph qualified as Gr
  import AtCoder.Internal.Assert qualified as ACIA
  import Data.Foldable (for_)
  import Data.Maybe (fromJust)
  import Data.Proxy (Proxy (..))
  import Data.Typeable (Typeable, typeRep)
  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 Test.QuickCheck.Classes qualified as QCC
  import Test.QuickCheck.Property qualified as QC
  import Test.Tasty
  import Test.Tasty.QuickCheck qualified as QC
  
  -- | Taken from `quickcheck-classes-base`.
  myForAllShrink ::
    (QC.Arbitrary a, Show b, Eq b) =>
    Bool -> -- Should we show the RHS. It's better not to show it if the RHS is equal to the input.
    (a -> Bool) -> -- is the value a valid input
    (a -> [String]) -> -- show the 'a' values
    String -> -- show the LHS
    (a -> b) -> -- the function that makes the LHS
    String -> -- show the RHS
    (a -> b) -> -- the function that makes the RHS
    QC.Property
  myForAllShrink displayRhs isValid showInputs name1 calc1 name2 calc2 =
    QC.MkProperty $
      QC.arbitrary >>= \x ->
        QC.unProperty $
          QC.shrinking QC.shrink x $ \x' ->
            let b1 = calc1 x'
                b2 = calc2 x'
                sb1 = show b1
                sb2 = show b2
                description = "  Description: " ++ name1 ++ " = " ++ name2
                err = description ++ "\n" ++ unlines (map ("  " ++) (showInputs x')) ++ "  " ++ name1 ++ " = " ++ sb1 ++ (if displayRhs then "\n  " ++ name2 ++ " = " ++ sb2 else "")
             in isValid x' QC.==> QC.counterexample err (b1 == b2)
  
  -- | Taken from `quickcheck-classes-base`.
  laws :: forall a. (Typeable a) => [Proxy a -> QCC.Laws] -> TestTree
  laws =
    testGroup (show (typeRep (Proxy @a)))
      . map
        ( \f ->
            let QCC.Laws name pairs = f (Proxy @a)
             in testGroup name (map (uncurry QC.testProperty) pairs)
        )
  
  -- | Returns an interval [l, r) in [0, n)
  intervalGen :: Int -> QC.Gen (Int, Int)
  intervalGen = intervalGen' 0
  
  -- | Returns an interval [l, r) in [0, n)
  intervalGen' :: Int -> Int -> QC.Gen (Int, Int)
  intervalGen' xl xr = do
    l <- QC.chooseInt (xl, xr)
    r <- QC.chooseInt (l, xr)
    pure (l, r)
  
  genDag :: forall w. (QC.Arbitrary w, VU.Unbox w) => Int -> QC.Gen (Gr.Csr w)
  genDag n = do
    edges <- VU.fromList <$> QC.sublistOf [(u, v) | u <- [0 .. n - 1], v <- [u + 1 .. n - 1]]
    verts <- VU.fromList <$> QC.shuffle [0 .. n - 1]
    ws <- VU.fromList <$> QC.vectorOf n (QC.arbitrary @w)
    let es = VU.zipWith (\(!u, !v) !w -> (verts VG.! u, verts VG.! v, w)) edges ws
    pure $ Gr.build n es
  
  -- | Newtype for generating a random DAG.
  newtype Dag w = Dag (Gr.Csr w)
  
  instance (QC.Arbitrary w, VU.Unbox w) => QC.Arbitrary (Dag w) where
    arbitrary = do
      QC.Positive n <- QC.arbitrary
      Dag <$> genDag n
  
  -- TODO: commit the upper code first
  
  -- | Prüfer sequence: https://en.wikipedia.org/wiki/Pr%C3%BCfer_sequence
  genTreeEdge :: Int -> QC.Gen (VU.Vector (Int, Int))
  genTreeEdge n = case n of
    1 -> pure VU.empty
    2 -> pure $ VU.singleton (0, 1)
    _ -> do
      prufer <- QC.vectorOf (n - 2) (QC.chooseInt (0, n - 1))
      pure $ VU.create $ do
        -- sum(deg) = n + (n - 2) at the beginning
        deg <- VUM.replicate n (1 :: Int)
        for_ prufer $ \v -> do
          VGM.modify deg (+ 1) v
  
        edges <- VUM.unsafeNew (n - 1)
  
        -- sum(deg) will be 2(n - 1) - 2(n - 2) = 2
        for_ (zip [0 :: Int ..] prufer) $ \(!i, !u) -> do
          -- NOTE: At the beginning, deg[u] > 1, and deg[i] is never decreased until iterated
          !v <- fromJust . VU.findIndex (== 1) <$> VU.unsafeFreeze deg
          -- So we're sure u /= v
          let !_ = ACIA.runtimeAssert (u /= v) "u /= v"
          VGM.write edges i (u, v)
          VGM.modify deg (subtract 1) u
          VGM.modify deg (subtract 1) v
  
        -- The last two vertices with degree one make up the last edge
        uv <- VU.findIndices (== 1) <$> VU.unsafeFreeze deg
        let !_ = ACIA.runtimeAssert (VU.length uv == 2) "not found uv?"
        VGM.write edges (n - 1) (uv VG.! 0, uv VG.! 1)
  
        -- done
        pure edges
  
  genTree :: forall w. (QC.Arbitrary w, VU.Unbox w) => Int -> QC.Gen (Gr.Csr w)
  genTree n = do
    es <- genTreeEdge n
    ws <- VU.fromList <$> QC.vectorOf (n - 1) (QC.arbitrary @w)
    pure . Gr.build n . Gr.swapDupe $ VU.zipWith (\(!u, !v) !w -> (u, v, w)) es ws
  
  -- | Newtype for generating a random tree.
  newtype Tree w = Tree (Gr.Csr w)
  
  instance (QC.Arbitrary w, VU.Unbox w) => QC.Arbitrary (Tree w) where
    arbitrary = do
      QC.Positive n <- QC.arbitrary
      Tree <$> genTree n
  