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(
        f"""
        MATCH (n)
        RETURN COUNT(n) AS count, LABEL(n) AS node_group
        """
    )
    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(
        f"""
        MATCH ()-[r]->()
        RETURN COUNT(r) AS count, LABEL(r) AS edge_group
        """
    )
    df: DataFrame = query_result.get_as_df()  # type: ignore
    return {row["edge_group"]: row["count"] for _, row in df.iterrows()}
