from __future__ import annotations

from typing import Any, Generator

import pandas as pd

from graph.connections.db_connection import DbConnection
from graph.exploration.algorithms.backend_service import (
    GraphAnalyticsExecutionLocation,
    GraphAnalyticsSamplingMode,
    get_graph_analytics_backend_service,
)
from graph.exploration.algorithms.graph_clustering_params import (
    GraphClusteringParams,
    selected_graph_clustering_names,
)
from graph.exploration.algorithms.projected_graph import ProjectedGraph


def compute_graph_clustering(
    conn: DbConnection,
    projected_graph: ProjectedGraph,
    params: GraphClusteringParams,
    batch_size: int = 10000,
    *,
    location: GraphAnalyticsExecutionLocation = GraphAnalyticsExecutionLocation.IN_MEMORY,
    sampling_mode: GraphAnalyticsSamplingMode | None = None,
    sampling_limit: int | None = None,
) -> Generator[pd.DataFrame, Any, None]:
    return get_graph_analytics_backend_service().execute(
        "graph_clustering",
        conn,
        projected_graph,
        params,
        batch_size=batch_size,
        location=location,
        sampling_mode=sampling_mode,
        sampling_limit=sampling_limit,
    )

__all__ = [
    "GraphClusteringParams",
    "ProjectedGraph",
    "compute_graph_clustering",
    "selected_graph_clustering_names",
]
