"""Visualization tools and graph algorithms for frequent itemsets."""

from __future__ import annotations

import typing

import pandas as pd

if typing.TYPE_CHECKING:
    import networkx


def to_networkx(
    rules_df: pd.DataFrame,
    source_col: str = "antecedents",
    target_col: str = "consequents",
    edge_attr: str = "lift",
) -> networkx.DiGraph:  # type: ignore
    """Converts a Rusket association rules DataFrame into a NetworkX Directed Graph.

    Nodes represent individual items. Directed edges represent rules (antecedent -> consequent).
    Edge weights are determined by the `edge_attr` parameter (typically lift or confidence).

    This is extremely useful for running community detection (e.g., Louvain)
    to automatically discover "Product Clusters" or visualising cross-selling patterns.

    Args:
        rules_df: A Pandas DataFrame generated by `rusket.association_rules()`.
        source_col: Column name containing antecedents.
        target_col: Column name containing consequents.
        edge_attr: The metric to use as the edge weight/thickness.

    Returns:
        networkx.DiGraph: A directed graph of the association rules.
    """
    try:
        import networkx as nx
    except ImportError as err:
        raise ImportError(
            "The 'networkx' library is required to use `to_networkx()`. "
            "Install it via `pip install networkx` or `uv add networkx`."
        ) from err

    if rules_df.empty:
        return nx.DiGraph()

    if source_col not in rules_df.columns or target_col not in rules_df.columns:
        raise ValueError(f"Columns {source_col} and {target_col} must exist in the DataFrame.")

    if edge_attr is not None and edge_attr not in rules_df.columns:
        raise ValueError(f"Edge attribute column '{edge_attr}' not found in the DataFrame.")

    G = nx.DiGraph()

    for _, row in rules_df.iterrows():
        # Rules can have multiple items in antecedents/consequents.
        # For a graph, we create edges from every item in antecedent to every item in consequent.
        # Note: If antecedents -> consequents is (A, B) -> (C),
        # we create A->C and B->C with the same rule weight.

        antecedents = list(row[source_col])
        consequents = list(row[target_col])
        weight = row[edge_attr] if edge_attr else 1.0

        for ant in antecedents:
            for con in consequents:
                if G.has_edge(ant, con):
                    # If multiple rules imply the same edge, keep the strongest one
                    if G[ant][con]["weight"] < weight:
                        G[ant][con]["weight"] = weight
                else:
                    G.add_edge(ant, con, weight=weight)

    return G
