content: module Tests.Extra.Graph where
  
  import AtCoder.Extra.Graph qualified as Gr
  import Control.Monad (unless)
  import Control.Monad.Fix (fix)
  import Data.List qualified as L
  import Data.Vector qualified as V
  import Data.Vector.Generic qualified as VG
  import Data.Vector.Unboxed qualified as VU
  import Data.Vector.Unboxed.Mutable qualified as VUM
  import Test.Tasty
  import Test.Tasty.HUnit
  import Test.Tasty.QuickCheck as QC
  import Tests.Util (genDag)
  
  reachableFlags :: Int -> (Int -> VU.Vector Int) -> Int -> VU.Vector Bool
  reachableFlags n gr u0 = VU.create $ do
    vis <- VUM.replicate n False
    VUM.write vis u0 True
    flip fix u0 $ \loop u -> do
      VU.forM_ (gr u) $ \v -> do
        b <- VUM.exchange vis v True
        unless b $ do
          loop v
    pure vis
  
  testTopSort :: Int -> Gr.Csr () -> VU.Vector Int -> Bool
  testTopSort n gr vs =
    let reachables = V.generate n (reachableFlags n (gr `Gr.adj`))
     in and
          [ not $ reachables VG.! v VG.! u
            | iu <- [0 .. n - 1],
              let u = vs VG.! iu,
              iv <- [iu + 1 .. n - 1],
              let v = vs VG.! iv
          ]
  
  -- | Tests lexicographically smallest topological ordering.
  prop_topSort :: QC.Gen QC.Property
  prop_topSort = do
    n <- QC.chooseInt (1, 3)
    dag <- genDag @() n
    let vs = Gr.topSort n (dag `Gr.adj`)
    let perms = map (VU.fromListN n) $ L.permutations [0 .. n - 1]
    pure $ vs QC.=== minimum (filter (testTopSort n dag) perms)
  
  genComplexEdges :: Int -> QC.Gen (VU.Vector (Int, Int, Int))
  genComplexEdges n = do
    m <- QC.chooseInt (1, 2 * n * n)
    (VU.fromList <$>) . QC.vectorOf m $ do
      u <- QC.chooseInt (0, n - 1)
      v <- QC.chooseInt (0, n - 1)
      w <- QC.arbitrary @Int
      pure (u, v, w)
  
  prop_floydWarshall :: QC.Gen QC.Property
  prop_floydWarshall = do
    -- n <- QC.chooseInt (1, 16)
    let n = 4
    es <- genComplexEdges n
    let !undefW = maxBound `div` 2 :: Int
    let (!distFw, !_prevFw) = Gr.trackingFloydWarshall n es undefW
    let gr = Gr.build n es
    let !bell = V.generate n $ Gr.trackingBellmanFord n (Gr.adjW gr) undefW . VU.singleton . (,0)
    pure $
      QC.counterexample (show (n, es)) $
        QC.conjoin
          [ case bell VG.! u of
              -- TODO: assertion function?
              Nothing -> any (\vtx -> distFw VG.! (n * vtx + vtx) < 0) [0 .. n - 1] QC.=== True
              Just (!distB, !_prevB) ->
                QC.conjoin
                  [ distFw VG.! (n * u + v) QC.=== distB VG.! v
                  -- TODO: Shortest paths cannot be uniqueified, so other test would be suitable
                  -- , Gr.constructPathFromRootNN prevFw u v QC.=== Gr.constructPathFromRoot prevB v
                  ]
            | u <- [0 .. n - 1],
              v <- [0 .. n - 1]
          ]
  
  unit_loopPathConstruction :: TestTree
  unit_loopPathConstruction = testCase "loop path reconstruction" $ do
    let parents = VU.fromList [3, 0, 1, 2]
    let path = Gr.constructPathFromRoot parents 3
    path @?= VU.fromList [0, 1, 2, 3]
  
  unit_one :: TestTree
  unit_one = testCase "one" $ do
    let !gr = Gr.build @Int 1 VU.empty
    Gr.topSort 1 (Gr.adj gr) @?= VU.singleton 0
    Gr.scc gr @?= V.singleton (VU.singleton 0)
    Gr.rev gr @?= gr
    Gr.findCycleDirected gr @?= Nothing
    Gr.findCycleUndirected gr @?= Nothing
    Gr.connectedComponents 1 (Gr.adj gr) @?= V.singleton (VU.singleton 0)
    (VU.length <$> (Gr.bipartiteVertexColors 1 (Gr.adj gr))) @?= Just 1
    -- Gr.blockCut 1 (Gr.adj gr)
    -- Gr.blockCutComponents 1 (Gr.adj gr)
    Gr.bfs 1 (Gr.adjW gr) (-1 :: Int) (VU.singleton (0, 7)) @?= (VU.singleton 7)
    Gr.trackingBfs 1 (Gr.adjW gr) (-1 :: Int) (VU.singleton (0, 7)) @?= (VU.singleton 7, VU.singleton (-1))
    Gr.bfs01 1 0 (Gr.adjW gr) (VU.singleton (0, 7)) @?= (VU.singleton 7)
    Gr.trackingBfs01 1 0 (Gr.adjW gr) (VU.singleton (0, 7)) @?= (VU.singleton 7, VU.singleton (-1))
    Gr.dijkstra 1 0 (Gr.adjW gr) (-1 :: Int) (VU.singleton (0, 7)) @?= (VU.singleton 7)
    Gr.trackingDijkstra 1 0 (Gr.adjW gr) (-1 :: Int) (VU.singleton (0, 7)) @?= (VU.singleton 7, VU.singleton (-1))
    Gr.bellmanFord 1 (Gr.adjW gr) (-1 :: Int) (VU.singleton (0, 7)) @?= Just (VU.singleton 7)
    Gr.trackingBellmanFord 1 (Gr.adjW gr) (-1 :: Int) (VU.singleton (0, 7)) @?= Just (VU.singleton 7, VU.singleton (-1))
    Gr.floydWarshall 1 VU.empty (maxBound `div` 2 :: Int) @?= VU.singleton 0
    Gr.trackingFloydWarshall 1 VU.empty (maxBound `div` 2 :: Int) @?= (VU.singleton 0, VU.singleton (-1))
    (!fw, !prev) <- Gr.newTrackingFloydWarshall 1 VU.empty (maxBound `div` 2 :: Int)
    let test = do
          fw' <- VU.unsafeFreeze fw
          prev' <- VU.unsafeFreeze prev
          (fw', prev') @?= (VU.singleton 0, VU.singleton (-1))
    test
    -- add positive self-loop edge
    Gr.updateEdgeFloydWarshall fw 1 (maxBound `div` 2) 0 0 1
    test
    -- add negative self-loop edge
    Gr.updateEdgeFloydWarshall fw 1 (maxBound `div` 2) 0 0 (-1)
    {- TODO: write test here -}
    pure ()
  
  tests :: [TestTree]
  tests =
    [ QC.testProperty "topSort" prop_topSort,
      -- not writing much tests, as we have verification problems
      QC.testProperty "floydWarshall" prop_floydWarshall,
      unit_one
    ]
  