from pandas import DataFrame

from graph.connections.db_connection import DbConnection
from graph.exploration.queries.total_counts import GroupCounts


def get_node_counts_by_group(conn: DbConnection) -> GroupCounts:
    query_result = conn.execute(
        """
        MATCH (n)
        WITH labels(n) AS node_labels, n
        UNWIND node_labels AS node_label
        RETURN node_label AS node_group, COUNT(n) AS count
        """
    )
    df: DataFrame = query_result.get_as_df()  # type: ignore
    return {row["node_group"]: row["count"] for _, row in df.iterrows()}


def get_edge_counts_by_group(conn: DbConnection) -> GroupCounts:
    query_result = conn.execute(
        """
        MATCH ()-[r]->()
        RETURN type(r) AS edge_group, COUNT(r) AS count
        """
    )
    df: DataFrame = query_result.get_as_df()  # type: ignore
    return {row["edge_group"]: row["count"] for _, row in df.iterrows()}
