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

from __future__ import annotations

import typing
from typing import Any

if typing.TYPE_CHECKING:
    import networkx  # type: ignore
    import numpy as np

    from rusket._dependencies import import_optional_dependency

    pd = import_optional_dependency("pandas")


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

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

    This is extremely useful for running community detection algorithms
    (e.g., Louvain, Girvan-Newman) to automatically discover **product clusters**,
    or for visualising cross-selling patterns as a force-directed graph.

    Parameters
    ----------
    rules_df : pd.DataFrame
        A Pandas DataFrame generated by ``rusket.association_rules()``.
    source_col : str, default='antecedents'
        Column name containing antecedents (graph edge sources).
    target_col : str, default='consequents'
        Column name containing consequents (graph edge targets).
    edge_attr : str, default='lift'
        The metric to use as edge weight/thickness.

    Returns
    -------
    networkx.DiGraph
        A directed graph of the association rules. If ``rules_df`` is empty,
        returns an empty ``DiGraph``.

    Notes
    -----
    Requires the ``networkx`` package (``pip install networkx``).
    When multiple rules produce the same directed edge, only the highest-weight
    rule is retained.

    Examples
    --------
    >>> import rusket
    >>> G = rusket.viz.to_networkx(rules_df, edge_attr="lift")
    >>> # Community detection with networkx
    >>> import networkx.algorithms.community as nx_comm
    >>> communities = nx_comm.greedy_modularity_communities(G.to_undirected())
    """
    try:
        import networkx as nx  # type: ignore
    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


def to_networkxr(
    rules: Any,
    source_col: str = "antecedents",
    target_col: str = "consequents",
    edge_attr: str = "lift",
) -> Any:
    """Convert Rusket association rules into a NetworkXR Directed Graph.

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

    This is extremely useful for running algorithms to automatically discover
    **product clusters**, or for visualising cross-selling patterns.

    Parameters
    ----------
    rules : Any
        A list of dictionaries, a Pandas DataFrame, or a Polars DataFrame
        generated by ``rusket.association_rules()``.
    source_col : str, default='antecedents'
        Column name containing antecedents (graph edge sources).
    target_col : str, default='consequents'
        Column name containing consequents (graph edge targets).
    edge_attr : str, default='lift'
        The metric to use as edge weight/thickness.

    Returns
    -------
    networkxr.DiGraph
        A directed graph of the association rules. If ``rules`` is empty,
        returns an empty ``DiGraph``.

    Notes
    -----
    Requires the ``networkxr`` package (``pip install rusket[network]``).
    When multiple rules produce the same directed edge, only the highest-weight
    rule is retained.

    Examples
    --------
    >>> import rusket
    >>> G = rusket.viz.to_networkxr(rules, edge_attr="lift")
    """
    try:
        import networkxr as nxr
    except ImportError as err:
        raise ImportError(
            "The 'networkxr' library is required to use `to_networkxr()`. "
            "Install it via `pip install networkxr` or `uv add networkxr`."
        ) from err

    records = rules
    if hasattr(rules, "to_dicts"):
        records = rules.to_dicts()
    elif hasattr(rules, "to_dict"):
        try:
            records = rules.to_dict(orient="records")
        except TypeError:
            records = rules.to_dict()
    elif hasattr(rules, "iterrows"):
        records = [row for _, row in rules.iterrows()]

    if not isinstance(records, (list, tuple)):
        records = list(records)

    if not records:
        return nxr.DiGraph()

    first = records[0]
    if source_col not in first or target_col not in first:
        raise ValueError(f"Columns {source_col} and {target_col} must exist in the data.")

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

    G = nxr.DiGraph()

    for row in records:
        antecedents = list(row[source_col])
        consequents = list(row[target_col])
        weight = float(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


def compute_pca_3d(data: np.ndarray) -> np.ndarray:
    """Compute 3D PCA coordinates for the given data.

    Uses the Rust-backed PCA implementation for performance.

    Parameters
    ----------
    data : np.ndarray
        Data matrix of shape (n_samples, n_features).

    Returns
    -------
    np.ndarray
        3D projection of shape (n_samples, 3).
    """
    from .pca import pca3

    return pca3(data)


def visualize_latent_space(
    model: Any,
    labels: bool = True,
    n_items: int | None = None,
) -> Any:
    """Visualizes the item latent space in 3D using PCA.

    Parameters
    ----------
    model : ImplicitRecommender
        A fitted model with `item_factors`.
    labels : bool, default=True
        Whether to show item labels on hover.
    n_items : int, optional
        Limit visualization to the first N items (useful for very large models).

    Returns
    -------
    plotly.graph_objects.Figure
        A Plotly figure object.
    """
    import numpy as np

    from rusket._dependencies import import_optional_dependency

    pd = import_optional_dependency("pandas")

    try:
        import plotly.express as px
    except ImportError as err:
        raise ImportError(
            "The 'plotly' library is required to use `visualize_latent_space()`. "
            "Install it via `pip install plotly` or `uv add plotly`."
        ) from err

    factors = model.item_factors
    if n_items is not None:
        factors = factors[:n_items]

    # Normalize factors for better visualization (cosine space)
    norms = np.linalg.norm(factors, axis=1, keepdims=True)
    factors_norm = factors / np.clip(norms, a_min=1e-10, a_max=None)

    pca_coords = compute_pca_3d(factors_norm)

    df_viz = pd.DataFrame(
        {
            "x": pca_coords[:, 0],
            "y": pca_coords[:, 1],
            "z": pca_coords[:, 2],
        }
    )

    if labels and hasattr(model, "_item_labels") and model._item_labels is not None:
        item_labels = model._item_labels
        if n_items is not None:
            item_labels = item_labels[:n_items]
        df_viz["item"] = item_labels
        hover_data = ["item"]
    else:
        hover_data = None

    fig = px.scatter_3d(
        df_viz,
        x="x",
        y="y",
        z="z",
        hover_data=hover_data,
        template="plotly_dark",
        title="Item Latent Space (3D PCA)",
    )

    # Tight layout and smaller markers for a premium look
    fig.update_traces(marker={"size": 3, "opacity": 0.8})
    fig.update_layout(
        margin={"l": 0, "r": 0, "b": 0, "t": 30},
        scene={
            "xaxis": {"showticklabels": False, "title": ""},
            "yaxis": {"showticklabels": False, "title": ""},
            "zaxis": {"showticklabels": False, "title": ""},
        },
    )

    return fig


def plot_pca(
    data: Any,
    labels: list[str] | None = None,
    title: str | None = None,
    marker_size: int = 4,
    opacity: float = 0.8,
    template: str = "plotly_dark",
) -> Any:
    """Create an interactive 2D or 3D scatter plot of PCA-projected data.

    Automatically detects whether to produce a 2D or 3D plot based on the
    number of columns in ``data``.

    Parameters
    ----------
    data : np.ndarray or ProjectedSpace
        Projected coordinates of shape ``(n_samples, 2)`` or ``(n_samples, 3)``,
        or a `ProjectedSpace` object returned by `pca()`.
    labels : list[str], optional
        Hover labels, one per sample.
    title : str, optional
        Plot title. Defaults to ``"PCA 2D"`` or ``"PCA 3D"``.
    marker_size : int, default=4
        Marker size.
    opacity : float, default=0.8
        Marker opacity.
    template : str, default="plotly_dark"
        Plotly template name.

    Returns
    -------
    plotly.graph_objects.Figure
        A Plotly figure object.  Call ``.show()`` to display it or
        ``.write_html()`` to save.

    Raises
    ------
    ImportError
        If ``plotly`` is not installed.
    ValueError
        If ``data`` does not have 2 or 3 columns.

    Examples
    --------
    >>> import numpy as np
    >>> import rusket
    >>> X = np.random.default_rng(0).standard_normal((200, 50)).astype("float32")
    >>> coords = rusket.pca3(X)
    >>> fig = rusket.viz.plot_pca(coords, title="My embedding")
    """
    try:
        import plotly.express as px  # type: ignore
    except ImportError as err:
        raise ImportError(
            "The 'plotly' library is required to use `plot_pca()`. "
            "Install it via `pip install plotly` or `uv add plotly`."
        ) from err

    import numpy as np

    from rusket._dependencies import import_optional_dependency

    pd = import_optional_dependency("pandas")

    if hasattr(data, "data"):
        data = data.data
    data = np.asarray(data)
    n_cols = data.shape[1] if data.ndim == 2 else 0

    if n_cols == 2:
        df = pd.DataFrame({"PC1": data[:, 0], "PC2": data[:, 1]})
        if labels is not None:
            df["label"] = labels
        fig = px.scatter(
            df,
            x="PC1",
            y="PC2",
            hover_data=["label"] if labels is not None else None,
            template=template,
            title=title or "PCA 2D",
        )
    elif n_cols == 3:
        df = pd.DataFrame({"PC1": data[:, 0], "PC2": data[:, 1], "PC3": data[:, 2]})
        if labels is not None:
            df["label"] = labels
        fig = px.scatter_3d(
            df,
            x="PC1",
            y="PC2",
            z="PC3",
            hover_data=["label"] if labels is not None else None,
            template=template,
            title=title or "PCA 3D",
        )
    else:
        raise ValueError(
            f"plot_pca expects data with 2 or 3 columns, got {n_cols}. "
            "Pass the output of rusket.pca2() or rusket.pca3()."
        )

    fig.update_traces(marker={"size": marker_size, "opacity": opacity})
    fig.update_layout(margin={"l": 20, "r": 20, "b": 20, "t": 40})

    return fig
