"""Similarity measures over AGE graphs via NetworkX."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from piggie._integrations.networkx import to_networkx
from piggie.exceptions import AlgorithmError

if TYPE_CHECKING:
    import pandas as pd

    from piggie.driver import PiggieDriver


def jaccard_similarity(
    driver: PiggieDriver,
    graph: str,
    node_pairs: list[tuple[int, int]] | None = None,
    *,
    label: str | None = None,
    rel_type: str | None = None,
) -> pd.DataFrame:
    """Compute Jaccard similarity coefficient for node pairs.

    Args:
        driver: Connected PiggieDriver.
        graph: Graph name.
        node_pairs: List of (u, v) tuples. None = all non-adjacent pairs.
        label: Filter to a single vertex label.
        rel_type: Filter to a single edge type.

    Returns:
        DataFrame with columns [node_u, node_v, score].
    """
    try:
        import networkx as nx
        import pandas as pd
    except ImportError as exc:
        raise ImportError(
            "networkx and pandas are required. "
            "Install with: pip install piggie[networkx,pandas]"
        ) from exc

    labels = [label] if label else None
    rel_types = [rel_type] if rel_type else None
    G = to_networkx(driver, graph, labels=labels, rel_types=rel_types)
    G_undirected = G.to_undirected()

    try:
        preds: Any = nx.jaccard_coefficient(G_undirected, ebunch=node_pairs)
        rows = [{"node_u": u, "node_v": v, "score": p} for u, v, p in preds]
    except Exception as exc:
        raise AlgorithmError(f"Jaccard similarity failed: {exc}") from exc

    return pd.DataFrame(rows, columns=["node_u", "node_v", "score"])


def adamic_adar(
    driver: PiggieDriver,
    graph: str,
    node_pairs: list[tuple[int, int]] | None = None,
    *,
    label: str | None = None,
    rel_type: str | None = None,
) -> pd.DataFrame:
    """Compute Adamic-Adar index for node pairs.

    Args:
        driver: Connected PiggieDriver.
        graph: Graph name.
        node_pairs: List of (u, v) tuples. None = all non-adjacent pairs.
        label: Filter to a single vertex label.
        rel_type: Filter to a single edge type.

    Returns:
        DataFrame with columns [node_u, node_v, score].
    """
    try:
        import networkx as nx
        import pandas as pd
    except ImportError as exc:
        raise ImportError(
            "networkx and pandas are required. "
            "Install with: pip install piggie[networkx,pandas]"
        ) from exc

    labels = [label] if label else None
    rel_types = [rel_type] if rel_type else None
    G = to_networkx(driver, graph, labels=labels, rel_types=rel_types)
    G_undirected = G.to_undirected()

    try:
        preds: Any = nx.adamic_adar_index(G_undirected, ebunch=node_pairs)
        rows = [{"node_u": u, "node_v": v, "score": p} for u, v, p in preds]
    except Exception as exc:
        raise AlgorithmError(f"Adamic-Adar computation failed: {exc}") from exc

    return pd.DataFrame(rows, columns=["node_u", "node_v", "score"])


def common_neighbors(
    driver: PiggieDriver,
    graph: str,
    u: int,
    v: int,
    *,
    label: str | None = None,
    rel_type: str | None = None,
) -> pd.DataFrame:
    """Find common neighbors of two nodes.

    Args:
        driver: Connected PiggieDriver.
        graph: Graph name.
        u: First node ID (AGE vertex id).
        v: Second node ID (AGE vertex id).
        label: Filter to a single vertex label.
        rel_type: Filter to a single edge type.

    Returns:
        DataFrame with columns [node_u, node_v, score] where score is the
        count of common neighbors.
    """
    try:
        import networkx as nx
        import pandas as pd
    except ImportError as exc:
        raise ImportError(
            "networkx and pandas are required. "
            "Install with: pip install piggie[networkx,pandas]"
        ) from exc

    labels = [label] if label else None
    rel_types = [rel_type] if rel_type else None
    G = to_networkx(driver, graph, labels=labels, rel_types=rel_types)
    G_undirected = G.to_undirected()

    try:
        neighbors = list(nx.common_neighbors(G_undirected, u, v))
    except Exception as exc:
        raise AlgorithmError(f"Common neighbors failed: {exc}") from exc

    return pd.DataFrame([{"node_u": u, "node_v": v, "score": len(neighbors)}])
