content: {-# LANGUAGE RecordWildCards #-}
  
  -- | Not complete!
  module Tests.Extra.Tree.Lct where
  
  import AtCoder.Extra.Graph qualified as Gr
  import AtCoder.Extra.Tree.Lct (Lct (..))
  import AtCoder.Extra.Tree.Lct qualified as Lct
  import AtCoder.Internal.Buffer qualified as B
  import Control.Monad (when)
  import Control.Monad.Primitive (PrimMonad, PrimState)
  import Data.Foldable (for_)
  import Data.Maybe (fromJust, isJust)
  import Data.Semigroup (Sum (..))
  import Data.Vector qualified as V
  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.Monadic as QCM
  import Test.Tasty
  import Test.Tasty.QuickCheck as QC
  
  -- | Maximum number of vertices.
  maxN :: Int
  maxN = 16
  
  -- | Maximum number of queries.
  maxQ :: Int
  maxQ = 1000
  
  -- | Range of monoid value.
  rng :: (Int, Int)
  rng = (-rngI, rngI)
  
  -- | Range of monoid value.
  rngI :: Int
  rngI = 10
  
  data Init s = Init
    { n :: {-# UNPACK #-} !Int,
      q :: {-# UNPACK #-} !Int,
      initialMonoids :: {-# UNPACK #-} !(VU.Vector (Sum Int)),
      refM :: !(IO Ref),
      lctM :: !(IO (Lct.Lct s (Sum Int)))
    }
  
  -- | Referencial implementation
  data Ref = Ref
    { nR :: {-# UNPACK #-} !Int,
      -- | Monoid values
      vecR :: !(VUM.MVector VUM.RealWorld (Sum Int)),
      -- | NxN
      edgeR :: !(VUM.MVector VUM.RealWorld Bool)
    }
  
  newRef :: Int -> VU.Vector (Sum Int) -> IO Ref
  newRef nR xs = do
    vecR <- VU.thaw xs
    edgeR <- VUM.replicate (nR * nR) False
    pure Ref {..}
  
  instance Show (Init s) where
    show Init {..} = show (n, q, initialMonoids)
  
  instance QC.Arbitrary (Init VUM.RealWorld) where
    arbitrary = do
      n <- QC.chooseInt (1, maxN)
      q <- QC.chooseInt (1, maxQ)
      initialMonoids <- (VU.fromList <$>) $ QC.vectorOf n (Sum <$> QC.chooseInt rng)
      let lctM = Lct.buildInv negate initialMonoids VU.empty
      let refM = let nR = n in newRef n initialMonoids
      pure Init {..}
  
  data Query
    = Read Int
    | Write Int (Sum Int)
    | Modify Int (Sum Int)
    | ModifyM Int (Sum Int)
    | Link Int Int
    | Cut Int Int
    | Evert Int
    | Expose Int
    | Expose_ Int
    | -- | Root Int
      Same Int Int
    | -- TODO: Test parent!
  
      -- | Parent Int
      -- | Jump
      JumpMaybe Int Int Int
    | LengthBetween Int Int
    | -- | Lca Int Int
      -- | LcaMaybe Int Int
      ProdPath Int Int
    | ProdSubtree Int Int
    -- -- | ProdTree Int Int
    deriving (Eq, Show)
  
  -- | Arbitrary return type for the `Query` result.
  data Result
    = None
    | I !Int
    | MI !(Maybe Int)
    | B Bool
    | SI (Sum Int)
    deriving (Show, Eq)
  
  queryGen :: Int -> QC.Gen Query
  queryGen n =
    QC.oneof
      [ Read <$> u,
        Write <$> u <*> x,
        Modify <$> u <*> x,
        ModifyM <$> u <*> x,
        Link <$> u <*> v,
        Cut <$> u <*> v,
        Evert <$> u,
        Expose <$> u,
        Expose_ <$> u,
        -- \| Root <$> Int
        Same <$> u <*> v,
        -- \| Parent <$> Int
        -- \| Jump
        JumpMaybe <$> u <*> v <*> k,
        LengthBetween <$> u <*> v,
        -- Lca <$> u <*> v,
        -- LcaMaybe <$> u <*> v,
        ProdPath <$> u <*> v
        -- TODO: test prod subtree
        -- ProdSubtree <$> u <*> v
      ]
    where
      u = QC.chooseInt (0, n - 1)
      v = QC.chooseInt (0, n - 1)
      x = Sum <$> QC.chooseInt rng
      k = QC.chooseInt (-1, n)
  
  toCsrM :: Ref -> IO (Gr.Csr ())
  toCsrM Ref {..} = do
    es <- VU.unsafeFreeze edgeR
    let uvs = (`VU.mapMaybe` VU.generate (nR * nR) id) $ \i ->
          let (!u, !v) = i `divMod` nR
           in if u /= v && es VG.! (nR * u + v)
                then Just (u, v)
                else Nothing
    pure $ Gr.build' nR uvs
  
  findPath :: Gr.Csr () -> Int -> Int -> Maybe (VU.Vector Int)
  findPath gr v1 v2
    | ws VG.! v2 == -1 = Nothing
    | otherwise = Just $ Gr.constructPathFromRoot tree v2
    where
      (!ws, !tree) = Gr.trackingBfs (Gr.nCsr gr) grF (-1 :: Int) src
      grF = Gr.adj1 gr
      src = VU.singleton (v1, 0)
  
  -- | containers. (referencial implementation)
  handleRef :: Ref -> Query -> IO Result
  handleRef ref@Ref {..} q = do
    case q of
      Read i -> do
        SI <$> VGM.read vecR i
      Write i x -> do
        VGM.write vecR i x
        pure None
      Modify i x -> do
        VGM.modify vecR (+ x) i
        pure None
      ModifyM i x -> do
        VGM.modifyM vecR (pure . (+ x)) i
        pure None
      Link u v -> do
        VGM.write edgeR (nR * u + v) True
        VGM.write edgeR (nR * v + u) True
        pure None
      Cut u v -> do
        VGM.write edgeR (nR * u + v) False
        VGM.write edgeR (nR * v + u) False
        pure None
      Evert _ -> do
        pure None
      Expose _ -> do
        pure None
      Expose_ _ -> do
        pure None
      -- \| Root Int
      Same u v -> do
        gr <- toCsrM ref
        pure . B . isJust $ findPath gr u v
      -- \| Parent Int
      -- \| Jump
      JumpMaybe u v k -> do
        gr <- toCsrM ref
        case findPath gr u v of
          Nothing -> pure $ MI Nothing
          Just path -> pure . MI $ path VG.!? k
      LengthBetween u v -> do
        gr <- toCsrM ref
        let !path = fromJust $ findPath gr u v
        pure . I $ VU.length path - 1
      -- LcaMaybe u v -> do
      --   gr <- toCsrM ref
      --   -- FIXME: HLD cannot be used
      --   let connected = isJust $ findPath gr u v
      --   if connected
      --     then pure . MI . Just $ Hld.lca (Hld.new gr) u v
      --     else pure $ MI Nothing
      ProdPath u v -> do
        gr <- toCsrM ref
        ws <- VU.unsafeFreeze vecR
        case findPath gr u v of
          Nothing -> pure $ SI mempty
          Just path -> pure . SI $ VU.foldMap (ws VG.!) path
      ProdSubtree u v -> do
        -- TODO: fix this
        pure None
  
  -- | ACL
  handleAcl :: (PrimMonad m) => Lct (PrimState m) (Sum Int) -> Query -> m Result
  handleAcl lct q = case q of
    Read i -> do
      SI <$> Lct.read lct i
    Write i x -> do
      Lct.write lct i x
      pure None
    Modify i x -> do
      Lct.modify lct (+ x) i
      pure None
    ModifyM i x -> do
      Lct.modifyM lct (pure . (+ x)) i
      pure None
    Link u v -> do
      -- TODO: detect if already connect
      Lct.link lct u v
      pure None
    Cut u v -> do
      Lct.cut lct u v
      pure None
    Evert i -> do
      Lct.evert lct i
      pure None
    Expose i -> do
      _ <- Lct.expose lct i
      pure None
    Expose_ i -> do
      Lct.expose_ lct i
      pure None
    -- \| Root Int
    Same u v -> do
      B <$> Lct.same lct u v
    -- \| Parent Int
    -- \| Jump
    JumpMaybe u v k -> do
      MI <$> Lct.jumpMaybe lct u v k
    LengthBetween u v -> do
      I <$> Lct.lengthBetween lct u v
    -- LcaMaybe u v -> do
    --   MI <$> Lct.lcaMaybe lct u v
    ProdPath u v -> do
      SI <$> Lct.prodPath lct u v
    ProdSubtree u v -> do
      SI <$> Lct.prodSubtree lct u v
  
  prop_randomTest :: Init VUM.RealWorld -> QC.Property
  prop_randomTest Init {..} = QCM.monadicIO $ do
    lct <- QCM.run lctM
    ref <- QCM.run refM
    qs <- QCM.pick $ QC.vectorOf q (queryGen n)
    record <- QCM.run $ B.new q
    for_ qs $ \query -> do
      b <- QCM.run $ filterRefM ref query
      when b $ do
        QCM.run $ B.pushBack record $ VU.DoNotUnboxStrict query
        record' <- QCM.run $ B.unsafeFreeze record
        expected <- QCM.run $ handleRef ref query
        result <- QCM.run $ handleAcl lct query
        when (expected /= result) $ do
          -- show all the run queries on failure
          let s = show . V.map (\(VU.DoNotUnboxStrict x) -> x) $ VU.convert record'
          QCM.monitor $ QC.counterexample $ "actual queries in use: " ++ s
        QCM.assertWith (expected == result) $ show (query, expected, result)
  
  -- | Filter invalid queries
  filterRefM :: Ref -> Query -> IO Bool
  filterRefM ref@Ref {..} q = case q of
    Read _ -> pure True
    Write _ _ -> pure True
    Modify _ _ -> pure True
    ModifyM _ _ -> pure True
    Link u v
      | u == v -> pure False
      | otherwise -> not <$> same u v
    Cut u v -> do
      es <- VU.unsafeFreeze edgeR
      pure $ es VG.! (nR * u + v)
    Evert _ -> pure True
    Expose _ -> pure True
    Expose_ _ -> pure True
    Same _ _ -> pure True
    JumpMaybe u v _ -> same u v
    LengthBetween u v -> same u v
    -- LcaMaybe u v -> same u v
    ProdPath u v -> same u v
    ProdSubtree u v -> same u v
    where
      same u v = do
        gr <- toCsrM ref
        let connected = isJust $ findPath gr u v
        pure connected
  
  tests :: [TestTree]
  tests =
    [ QC.testProperty "randomTest" prop_randomTest
    ]
  