{-# LANGUAGE RecordWildCards #-}

-- | Heavy-light decomposition is a method for partitioning a tree into segments with consecutive
-- indices. It processes various path queries in \(O(\log n)\) time. For segment tree integration
-- and monoid products, refer to the @TreeMonoid@ module.
--
-- ==== __Overview of the internals__
-- The following is for understanding the internals, not for using the API. Skip to the examples if
-- you want.
--
-- ===== Original tree
--
-- Consider a tree with arbitrary vertex order:
--
-- @
--  0--8--7--3--1--2--12--13--15--14     XX: original Vertex
--     |        |                        --: edge
-- 10--5        11--8--6                 |: edge
--     |
--     4
-- @
--
-- ===== `indexHld`: Vertex -> VertexHld
--
-- The tree vertices are reindexed with `indexHld`, where each segment is assigned consecutive
-- vertex indices:
--
-- @
--  0==1==2==3==4==5==6==7==8==9     XX: VertexHld
--     |        |                    ==: edges on the same segment
-- 14==13       10==11==12           |: edge between different segments
--     |
--     15
-- @
--
-- Note that vertices on higher (closer to the root) segments are assigned smaller indices. This is
-- internally very important when calculating `lca`.
--
-- ===== `headHld`: Vertex -> Vertex
--
-- `headHld` points the "head" vertex of each segment. It can be used for finding LCA of two
-- vertices. To find the LCA, move up to the head, go up to the parental segment's vertex and
-- repeat until the two vertices are on the same segment.
--
-- @
--  0==0==0==0==0==0==0==0==0==0     XX: original Vertex
--     |     |
--  5==5     11==11==11
--     |
--     4
-- @
--
-- `headHld` also works for identifying segments. When two vertices are on the same segment, they
-- have the same head.
--
-- ===== `parentHld`: Vertex -> Vertex
--
-- `parentHld` points the parental segment's vertex from a head:
--
-- @
-- (-1)==0==8==7==3==1==2==12==13==15     XX: original Vertex
--       |        |
--    5==8        1==11==8
--       |
--       5
-- @
--
-- ==== __Example__
-- Create an `Hld` for a tree:
--
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import AtCoder.Extra.Tree.Hld qualified as Hld
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> -- 0--1--2--3
-- >>> --    +
-- >>> --    +--4--5
-- >>> let n = 6
-- >>> -- note that the edges must be bi-directed:
-- >>> let tree = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (2, 3), (1, 4), (4, 5)]
-- >>> let hld = Hld.new tree
--
-- `Hld` can process various queries in \(O(\log n)\) time:
--
-- >>> Hld.ancestor hld 5 3 -- go up three parents from `5`
-- 0
--
-- >>> Hld.lca hld 2 5 -- lowest common ancestor of `2` and `5`:
-- 1
--
-- >>> Hld.jump hld 5 2 3   -- go to the third vertex from `5` to `2`:
-- Just 2
--
-- >>> Hld.lengthBetween hld 5 3 -- get the length (the number of edges) between `5` and `3`:
-- 4
--
-- >>> Hld.path hld 5 3     -- get the path between `5` and `3`:
-- [5,4,1,2,3]
--
-- So the subtree queries are run with recpect to root vertex `0`, since our HLD is created with
-- `new`:
--
-- >>> Hld.isInSubtree hld 2 3 -- `3` is in the subtree of `2`
-- True
--
-- >>> Hld.isInSubtree hld 2 4 -- `4` is not in the subtree of `2`
-- False
--
-- ===== Segment queries
-- Products and segment queries are primarily used by the @TreeMonoid@ module and is not intended
-- for direct use, but here's some examples. This time the reindex by the HLD is identity:
--
-- >>> Hld.indexHld hld
-- [0,1,2,3,4,5]
--
-- So we can easily understand the outputs:
--
-- >>> Hld.pathSegmentsInclusive Hld.WeightsAreOnVertices hld 5 3
-- [(5,4),(1,3)]
--
-- >>> Hld.pathSegmentsInclusive Hld.WeightsAreOnEdges hld 5 3 -- LCA (1) is removed
-- [(5,4),(2,3)]
--
-- >>> Hld.subtreeSegmentInclusive hld 1
-- (1,5)
--
-- @since 1.1.0.0
module AtCoder.Extra.Tree.Hld
  ( -- * Hld
    Hld (..),
    Vertex,
    VertexHld,

    -- * Constructors
    new,
    newAt,

    -- * LCA
    lca,

    -- * Jump
    ancestor,
    jump,

    -- * Path
    lengthBetween,
    path,
    pathSegmentsInclusive,

    -- * Subtree
    subtreeSegmentInclusive,
    isInSubtree,

    -- * Products
    WeightPolicy (..),
    prod,
  )
where

import AtCoder.Extra.Graph qualified as Gr
import AtCoder.Internal.Assert qualified as ACIA
import Control.Monad
import Control.Monad.Fix
import Control.Monad.ST
import Data.Maybe
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 GHC.Stack (HasCallStack)

-- | Original graph vertex.
--
-- @since 1.1.0.0
type Vertex = Int

-- | Vertex reindexed by `indexHld`. It's contiguous for each segment.
--
-- @since 1.1.0.0
type VertexHld = Vertex

-- | `Hld` partitions a tree into segments and assigns contiguous `VertexHld` for each segment.
--
-- @since 1.1.0.0
data Hld = Hld
  { -- | The root vertex.
    --
    -- @since 1.1.0.0
    rootHld :: {-# UNPACK #-} !Vertex,
    -- | Maps `Vertex` to the parent `Vertex`. Returns @-1@ for the root node.
    --
    -- @since 1.1.0.0
    parentHld :: !(VU.Vector Vertex),
    -- | Maps `Vertex` to `VertexHld`, re-indexed vertices contiguous in each segment.
    --
    -- @since 1.1.0.0
    indexHld :: !(VU.Vector VertexHld),
    -- | Maps `Vertex` to the head `Vertex` of the segment.
    --
    -- @since 1.1.0.0
    headHld :: !(VU.Vector Vertex),
    -- | Maps `VertexHld` back to `Vertex`. Used for `ancestor` etc.
    --
    -- @since 1.1.0.0
    revIndexHld :: !(VU.Vector Vertex),
    -- | Maps `Vertex` to their depth from the root. Used for `jump` etc.
    --
    -- @since 1.1.0.0
    depthHld :: !(VU.Vector Int),
    -- | Maps `Vertex` to the subtree size. This is for subtree products.
    --
    -- @since 1.1.0.0
    subtreeSizeHld :: !(VU.Vector Int)
  }
  deriving
    ( -- | @since 1.1.0.0
      Show,
      -- | @since 1.1.0.0
      Eq
    )

-- | \(O(n)\) Creates an `Hld` with \(0\) as the root vertex.
--
-- ==== Constraints
-- - \(n \ge 1\)
-- - The input graph must be a tree. Note that the edges must be undirected (both \((u, v, w)\) and
--   \((v, u, w)\) edges are required).
--
-- @since 1.1.0.0
{-# INLINE new #-}
new :: forall w. (HasCallStack) => Gr.Csr w -> Hld
new tree = newAt tree 0

-- | \(O(n)\) Creates an `Hld` with a root vertex specified.
--
-- ==== Constraints
-- - \(n \ge 1\)
-- - The input graph must be a tree. Note that the edges must be undirected (both \((u, v, w)\) and
--   \((v, u, w)\) edges are required).
--
-- @since 1.1.0.0
{-# INLINEABLE newAt #-}
newAt :: forall w. (HasCallStack) => Gr.Csr w -> Vertex -> Hld
newAt tree root
  | n == 1 =
      Hld
        0
        (VU.singleton (-1))
        (VU.singleton 0)
        (VU.singleton 0)
        (VU.singleton 0)
        (VU.singleton 0)
        (VU.singleton 1)
  | otherwise = runST $ do
      -- Re-create adjacent vertices so that the biggest subtree's head vertex comes first.
      --
      -- We /could/ instead record the biggest adjacent subtree vertex for each vertex, but the other
      -- DFS would be harder.
      let (!tree', !parent, !depths, !subtreeSize) = runST $ do
            adjVec <- VU.thaw (Gr.adjCsr tree)
            parent_ <- VUM.unsafeNew n
            depths_ <- VUM.unsafeNew n
            subtreeSize_ <- VUM.unsafeNew n

            _ <- (\f -> fix f 0 (-1) root) $ \loop depth p v1 -> do
              VGM.write parent_ v1 p
              VGM.write depths_ v1 depth

              (!size1, !eBig) <-
                VU.foldM'
                  ( \(!size1, !eBig) (!e2, !v2) -> do
                      if v2 == p
                        then pure (size1, eBig)
                        else do
                          size2 <- loop (depth + 1) v1 v2
                          -- NOTE: It's `>` because we should swap at least once if there's some vertex other
                          -- that the parent_.
                          pure (size1 + size2, if size1 > size2 then eBig else e2)
                  )
                  (1 :: Int, -1)
                  (tree `Gr.eAdj` v1)

              -- move the biggest subtree's head to the first adjacent vertex.
              -- it means the "heavy edge" or the longest segment.
              when (eBig /= -1) $ do
                VGM.swap adjVec eBig $ fst (VG.head (tree `Gr.eAdj` v1))

              -- record subtree size
              VGM.write subtreeSize_ v1 size1

              pure size1

            !vec <- VU.unsafeFreeze adjVec
            (tree {Gr.adjCsr = vec},,,)
              <$> VU.unsafeFreeze parent_
              <*> VU.unsafeFreeze depths_
              <*> VU.unsafeFreeze subtreeSize_

      -- vertex -> reindexed vertex index
      indices <- VUM.replicate n (-1 :: Int)

      -- vertex -> head vertex of the segment
      heads <- VUM.replicate n (-1 :: Int)

      _ <- (\f -> fix f (0 :: Int) root (-1) root) $ \loop acc h p v1 -> do
        -- reindex:
        VGM.write indices v1 acc
        let !acc' = acc + 1

        VGM.write heads v1 h

        -- when the first vertex is within the same segment:
        let (!adj1, !rest) = fromJust $ VU.uncons (tree' `Gr.adj` v1)
        acc'' <-
          if adj1 == p
            then pure acc'
            else loop acc' h v1 adj1

        -- the others are in other segments:
        VU.foldM'
          ( \a v2 -> do
              if v2 == p
                then pure a
                else loop a v2 v1 v2
          )
          acc''
          rest

      !indices' <- VU.unsafeFreeze indices
      let !revIndex = VU.update (VU.replicate n (-1)) $ VU.imap (flip (,)) indices'

      Hld root parent indices'
        <$> VU.unsafeFreeze heads
        <*> pure revIndex
        <*> pure depths
        <*> pure subtreeSize
  where
    !n = Gr.nCsr tree
    !_ = ACIA.runtimeAssert (2 * (Gr.nCsr tree - 1) == Gr.mCsr tree) "AtCoder.Extra.Hld.newAt: not an undirected tree"
    !_ = ACIA.runtimeAssert (n >= 1) "AtCoder.Extra.Hld.newAt: the tree must have at least one vertex"

-- | \(O(\log n)\) Calculates the lowest common ancestor of \(u\) and \(v\).
--
-- @since 1.1.0.0
{-# INLINEABLE lca #-}
lca :: (HasCallStack) => Hld -> Vertex -> Vertex -> Vertex
lca Hld {..} = inner
  where
    inner !x !y
      -- sort for easier processing
      -- TODO: @case compare ix iy@ would be easier for me to understand
      | ix > iy = inner y x
      -- @x@ and @y@ are in other segments:
      | hx /= hy = inner x $ parentHld VG.! hy
      -- @x@ and @y@ are within the same segment:
      -- select the smaller one, which is closer to the root and that is the LCA.
      | otherwise = x
      where
        !ix = indexHld VG.! x
        !iy = indexHld VG.! y
        hx = headHld VG.! x
        hy = headHld VG.! y

-- | \(O(\log n)\) Go up \(k\) times from a vertex \(v\) to the root node. Throws an error if \(k\)
-- is bigger than the depth of \(v\).
--
-- @since 1.1.0.0
{-# INLINEABLE ancestor #-}
ancestor :: (HasCallStack) => Hld -> Vertex -> Int -> Vertex
ancestor Hld {..} parent k0 = inner parent k0
  where
    !_ = ACIA.runtimeAssert (0 <= k0 && k0 <= depthHld VG.! parent) $ "AtCoder.Extra.Tree.Hld.ancestor: k-th ancestor is out of bounds (`k = " ++ show k0 ++ "`)"
    inner v k
      -- on this segment
      | k <= iv - ihv = revIndexHld VG.! (iv - k)
      -- next segment
      | otherwise = inner (parentHld VG.! hv) (k - (iv - ihv + 1))
      where
        iv = indexHld VG.! v
        hv = headHld VG.! v
        ihv = indexHld VG.! hv

-- | \(O(\log n)\) Returns the \(k\)-th vertex of the path between \(u\) and \(v\) from \(u\).
-- Returns `Nothing` if \(k\) is bigger than the path length or \(k < 0\).
--
-- @since 1.4.0.0
{-# INLINEABLE jump #-}
jump :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int -> Maybe Vertex
jump hld@Hld {..} u v k
  | k < 0 = Nothing
  | k > lenU + lenV = Nothing
  | k <= lenU = Just $ ancestor hld u k
  | otherwise = Just $ ancestor hld v (lenU + lenV - k)
  where
    lca_ = lca hld u v
    du = depthHld VG.! u
    dv = depthHld VG.! v
    lenU = du - depthHld VG.! lca_
    lenV = dv - depthHld VG.! lca_

-- | \(O(\log n)\) Returns the length of the path between \(u\) and \(v\).
--
-- @since 1.1.0.0
{-# INLINEABLE lengthBetween #-}
lengthBetween :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int
lengthBetween hld@Hld {..} u v = du - dLca + dv - dLca
  where
    !lca_ = lca hld u v
    !dLca = depthHld VG.! lca_
    !du = depthHld VG.! u
    !dv = depthHld VG.! v

-- | \(O(n)\) Returns the vertices on the path between \(u\) and \(v\).
--
-- @since 1.1.0.0
{-# INLINEABLE path #-}
path :: (HasCallStack) => Hld -> Vertex -> Vertex -> [Vertex]
path hld@Hld {..} u v = concatMap expand $ pathSegmentsInclusive WeightsAreOnVertices hld u v
  where
    expand (!l, !r)
      | l <= r = map (revIndexHld VG.!) [l .. r]
      | otherwise = map (revIndexHld VG.!) [l, l - 1 .. r]

-- | \(O(\log n)\) Decomposes a path between two vertices \(u\) and \(v\) into segments. Each
-- segment is represented as an __inclusive__ range \([u_i, v_i]\) of `VertexHLD`.
--
-- The LCA is omitted from the returning vertices when the weight policy is set to
-- `WeightsAreOnEdges`. This is the trick to put edge weights to on vertices.
--
-- @since 1.1.0.0
{-# INLINEABLE pathSegmentsInclusive #-}
pathSegmentsInclusive :: (HasCallStack) => WeightPolicy -> Hld -> Vertex -> Vertex -> [(VertexHld, VertexHld)]
pathSegmentsInclusive weightPolicy Hld {..} x0 y0 = done $ inner x0 [] y0 []
  where
    isEdge = weightPolicy == WeightsAreOnEdges
    done (!up, !down) = reverse up ++ down
    -- @up@: bottom to top. [(max, min)]
    -- @down@: top to bottom. [(min, max)]
    inner :: Vertex -> [(VertexHld, VertexHld)] -> Vertex -> [(VertexHld, VertexHld)] -> ([(VertexHld, VertexHld)], [(VertexHld, VertexHld)])
    inner x up y down
      | hx == hy && isEdge = case compare ix iy of
          -- skip LCA on edge vertices
          LT -> (up, (ix {- edge -} + 1, iy) : down)
          GT -> ((ix, iy {- edge -} + 1) : up, down)
          EQ -> (up, down)
      | hx == hy && not isEdge = case compare ix iy of
          LT -> (up, (ix, iy) : down)
          _ -> ((ix, iy) : up, down)
      | otherwise = case compare ix iy of
          LT -> inner x up phy ((ihy, iy) : down)
          GT -> inner phx ((ix, ihx) : up) y down
          EQ -> error "unreachable"
      where
        ix, iy :: VertexHld
        !ix = indexHld VG.! x
        !iy = indexHld VG.! y
        hx, hy :: Vertex
        hx = headHld VG.! x
        hy = headHld VG.! y
        ihx, ihy :: VertexHld
        ihx = indexHld VG.! hx
        ihy = indexHld VG.! hy
        phx, phy :: VertexHld
        phx = parentHld VG.! hx
        phy = parentHld VG.! hy

-- | \(O(1)\) Returns a half-open interval of `VertexHld` \([\mathrm{start}, \mathrm{end})\) that
-- corresponds to the subtree vertices rooted at the given vertex.
--
-- @since 1.1.0.0
{-# INLINEABLE subtreeSegmentInclusive #-}
subtreeSegmentInclusive :: (HasCallStack) => Hld -> Vertex -> (VertexHld, VertexHld)
subtreeSegmentInclusive Hld {..} subtreeRoot = (ir, ir + sr - 1)
  where
    ir = indexHld VG.! subtreeRoot
    sr = subtreeSizeHld VG.! subtreeRoot

-- | \(O(1)\) Returns `True` if \(u\) is a vertex of subtree rooted at \(r\).
--
-- @since 1.1.0.0
{-# INLINEABLE isInSubtree #-}
isInSubtree :: (HasCallStack) => Hld -> Vertex -> Vertex -> Bool
isInSubtree hld@Hld {..} r_ u = l <= iu && iu <= r
  where
    (!l, !r) = subtreeSegmentInclusive hld r_
    !iu = indexHld VG.! u

-- | Represents whether weights are put on vertices or edges.
--
-- @since 1.1.0.0
data WeightPolicy
  = -- | Weights are put on vertices.
    --
    -- @since 1.1.0.0
    WeightsAreOnVertices
  | -- | Weights are put on edges.
    --
    -- @since 1.1.0.0
    WeightsAreOnEdges
  deriving
    ( -- | @since 1.1.0.0
      Eq,
      -- | @since 1.1.0.0
      Show
    )

-- | \(O(\log n f)\) Returns product of the path between \(u\) and \(v\), using the user functions
-- of time complexity \(O(f)\).
--
-- @since 1.1.0.0
{-# INLINE prod #-}
prod ::
  (HasCallStack, Monoid mono, Monad m) =>
  -- | The `WeightPolicy`.
  WeightPolicy ->
  -- | The `Hld`.
  Hld ->
  -- | User function for getting products in \([u, v)\), where \(u < v\) and
  -- \(\mathrm{depth}(u) < \mathrm{depth}(v)\).
  (VertexHld -> VertexHld -> m mono) ->
  -- | User function for getting products in \([u, v)\), where \(u < v\) and
  -- \(\mathrm{depth}(u) > \mathrm{depth}(v)\).
  (VertexHld -> VertexHld -> m mono) ->
  -- | \(u\).
  Vertex ->
  -- | \(v\).
  Vertex ->
  -- | Product of the path between \(u\) and \(v\).
  m mono
prod weightPolicy hld prodF prodB u0 v0 = do
  foldM
    ( \ !acc (!u, !v) -> do
        !x <-
          if u <= v
            then prodF u $ v + 1
            else prodB v $ u + 1
        pure $! acc <> x
    )
    mempty
    (pathSegmentsInclusive weightPolicy hld u0 v0)