"""Graph embedding algorithms over AGE graphs."""

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 node2vec_embed(
    driver: PiggieDriver,
    graph: str,
    *,
    dimensions: int = 64,
    walk_length: int = 30,
    num_walks: int = 200,
    p: float = 1.0,
    q: float = 1.0,
    workers: int = 1,
    label: str | None = None,
    rel_type: str | None = None,
) -> pd.DataFrame:
    """Compute node2vec embeddings for nodes in an AGE graph.

    Requires the ``node2vec`` package. Install with::

        pip install piggie[embeddings]

    Args:
        driver: Connected PiggieDriver.
        graph: Graph name.
        dimensions: Embedding dimensionality (default 64).
        walk_length: Length of each random walk (default 30).
        num_walks: Number of random walks per node (default 200).
        p: Return parameter (default 1.0).
        q: In-out parameter (default 1.0).
        workers: Number of parallel workers for word2vec (default 1).
        label: Filter to a single vertex label.
        rel_type: Filter to a single edge type.

    Returns:
        DataFrame with columns [node_id, label, embedding].
        embedding is a list of floats with length ``dimensions``.
    """
    try:
        import pandas as pd
    except ImportError as exc:
        raise ImportError(
            "pandas is required. Install with: pip install piggie[pandas]"
        ) from exc

    try:
        from node2vec import Node2Vec
    except ImportError:
        raise ImportError(
            "node2vec is required for graph embeddings. "
            "Install with: pip install piggie[embeddings]"
        ) from None

    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)

    try:
        model: Any = Node2Vec(
            G,
            dimensions=dimensions,
            walk_length=walk_length,
            num_walks=num_walks,
            p=p,
            q=q,
            workers=workers,
            quiet=True,
        ).fit(window=10, min_count=1, batch_words=4)
    except Exception as exc:
        raise AlgorithmError(f"node2vec embedding failed: {exc}") from exc

    rows = []
    for node_id in G.nodes():
        node_label = G.nodes[node_id].get("label", "")
        try:
            embedding = model.wv[str(node_id)].tolist()
        except KeyError:
            embedding = [0.0] * dimensions
        rows.append({"node_id": node_id, "label": node_label, "embedding": embedding})

    return pd.DataFrame(rows)
