from __future__ import annotations

from collections import Counter
from itertools import combinations, groupby
from typing import Any

import pandas as pd


def compute_projected_graph(
    df: pd.DataFrame,
    *,
    graph_of_column: str,
    linked_by_column: str,
    weighted: bool,
) -> pd.DataFrame:
    projection_columns = [graph_of_column, linked_by_column]
    filtered_df = df.loc[
        df[graph_of_column].notna() & df[linked_by_column].notna(),
        projection_columns,
    ]
    deduplicated_df = filtered_df.drop_duplicates(subset=projection_columns)

    projected_pairs = _build_projected_pairs(
        deduplicated_df,
        graph_of_column=graph_of_column,
        linked_by_column=linked_by_column,
    )

    if weighted:
        output_df = pd.DataFrame(
            [(source, target, weight) for (source, target), weight in projected_pairs.items()],
            columns=[f"{graph_of_column}_1", f"{graph_of_column}_2", "weight"],
        )
        if not output_df.empty:
            output_df = output_df.sort_values(
                by=["weight", f"{graph_of_column}_1", f"{graph_of_column}_2"],
                ascending=[False, True, True],
            ).reset_index(drop=True)
    else:
        output_df = pd.DataFrame(
            list(projected_pairs),
            columns=[f"{graph_of_column}_1", f"{graph_of_column}_2"],
        )
        if not output_df.empty:
            output_df = output_df.sort_values(
                by=[f"{graph_of_column}_1", f"{graph_of_column}_2"],
                ascending=[True, True],
            ).reset_index(drop=True)

    return output_df


def _canonical_sort_key(node: Any) -> tuple[str, str]:
    return str(type(node)), str(node)


def _iter_projected_pairs(
    deduplicated_df: pd.DataFrame,
    *,
    graph_of_column: str,
    linked_by_column: str,
):
    # Build the projection directly from grouped pairs instead of materializing a full
    # bipartite NetworkX graph, which keeps the recipe's memory overhead lower. We sort
    # each linked-by group once so itertools.combinations yields canonical edge pairs
    # without a Python-level comparison on every generated pair.
    sorted_pairs = deduplicated_df.sort_values(
        by=[linked_by_column],
        kind="stable",
    )
    rows = sorted_pairs[[graph_of_column, linked_by_column]].itertuples(index=False, name=None)

    for _, group_rows in groupby(rows, key=lambda row: row[1]):
        grouped_nodes = [row[0] for row in group_rows]
        if len(grouped_nodes) < 2:
            continue

        grouped_nodes.sort(key=_canonical_sort_key)
        yield from combinations(grouped_nodes, 2)


def _build_projected_pairs(
    deduplicated_df: pd.DataFrame,
    *,
    graph_of_column: str,
    linked_by_column: str,
) -> Counter[tuple[Any, Any]]:
    return Counter(
        _iter_projected_pairs(
            deduplicated_df,
            graph_of_column=graph_of_column,
            linked_by_column=linked_by_column,
        )
    )
