<gh_stars>0
from kgmk.dsa.topology.graph.jit.sort_csgraph import (
  sort_csgraph,
)

# TODO cut below
import numpy as np 
import numba as nb 


@nb.njit
def strongly_connected_components(
  n: int, 
  g: np.ndarray,
) -> np.ndarray:
  def scc_dfs(n, g): 
    g, edge_idx, _ = sort_csgraph(n, g)
    que = np.empty(n, np.int64)
    ptr = -1 
    visited = np.zeros(n, np.bool8)
    for i in range(n):
      if visited[i]: continue
      st = [i]
      while st:
        u = st.pop()
        if u < 0:
          que[ptr] = ~u
          ptr -= 1
          continue
        if visited[u]: continue
        visited[u] = True
        st.append(~u)
        for v in g[edge_idx[u]:edge_idx[u + 1], 1][::-1]:
          if not visited[v]: st.append(v)
    return que

  def scc_reverse_dfs(n, g, que):
    g[:, :2] = g[:, 1::-1]
    g, edge_idx, _ = sort_csgraph(n, g)
    label = np.full(n, -1, np.int64)
    l = 0
    for i in que:
      if label[i] != -1: continue
      st = [i]
      label[i] = l
      while st:
        u = st.pop()
        for v in g[edge_idx[u]:edge_idx[u + 1], 1][::-1]:
          if label[v] != -1: continue
          label[v] = l
          st.append(v)
      l += 1
    return label
  
  g = g.copy()
  que = scc_dfs(n, g)
  return scc_reverse_dfs(n, g, que)