"""Visualize COGILES colored graphs using matplotlib and networkx."""

from __future__ import annotations

import networkx as nx

from cogiles.parser import parse

# Layout algorithms mapped by name.
LAYOUT_MAP: dict[str, object] = {
    "kamada_kawai": nx.kamada_kawai_layout,
    "spring": nx.spring_layout,
    "circular": nx.circular_layout,
    "shell": nx.shell_layout,
    "spectral": nx.spectral_layout,
    "planar": nx.planar_layout,
    "random": nx.random_layout,
}


def _import_matplotlib():
    """Lazy-import matplotlib.pyplot, raising a helpful error if missing."""
    try:
        import matplotlib.pyplot as plt

        return plt
    except ImportError:
        raise ImportError(
            "matplotlib is required for visualization. "
            "Install it with: pip install cogiles[viz]"
        ) from None


def _label_color(rgb: tuple[float, float, float]) -> str:
    """Return 'white' or 'black' for maximum contrast against *rgb*."""
    luminance = 0.299 * rgb[0] + 0.587 * rgb[1] + 0.114 * rgb[2]
    return "white" if luminance < 0.5 else "black"


def draw(
    graph: nx.Graph | str,
    *,
    ax=None,
    layout: str = "kamada_kawai",
    labels: str | None = None,
    node_size: int = 300,
    edge_color: str = "black",
    node_edge_color: str = "black",
    node_edge_width: float = 1.0,
):
    """Draw a COGILES colored graph.

    Parameters
    ----------
    graph : nx.Graph or str
        A NetworkX graph with COGILES node attributes, or a COGILES string
        that will be parsed automatically.
    ax : matplotlib.axes.Axes, optional
        Axes to draw on.  If *None*, a new figure and axes are created.
    layout : str, optional
        Layout algorithm name.  One of ``"kamada_kawai"`` (default),
        ``"spring"``, ``"circular"``, ``"shell"``, ``"spectral"``,
        ``"planar"``, ``"random"``.
    labels : str or None, optional
        Node label style: ``None`` (no labels, default), ``"symbol"``
        (color symbol like ``"R"``), or ``"index"`` (node index).
    node_size : int, optional
        Size of the drawn nodes (default 300).
    edge_color : str, optional
        Color of the edges (default ``"black"``).
    node_edge_color : str, optional
        Outline/border color for nodes (default ``"black"``).
    node_edge_width : float, optional
        Width of the node outline (default 1.0).

    Returns
    -------
    tuple[Figure, Axes]
        The matplotlib Figure and Axes used for drawing.

    Raises
    ------
    CogilesParseError
        If *graph* is a string and cannot be parsed.
    ValueError
        If *layout* or *labels* is not a recognized value.
    ImportError
        If matplotlib is not installed.
    """
    if isinstance(graph, str):
        graph = parse(graph)

    plt = _import_matplotlib()

    if layout not in LAYOUT_MAP:
        valid = ", ".join(sorted(LAYOUT_MAP))
        raise ValueError(f"Unknown layout {layout!r}. Choose from: {valid}")

    if labels not in (None, "symbol", "index"):
        raise ValueError(
            f"Unknown labels {labels!r}. Choose from: None, 'symbol', 'index'"
        )

    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = ax.get_figure()

    pos = LAYOUT_MAP[layout](graph)

    node_colors = [graph.nodes[n]["color"] for n in graph.nodes()]

    nx.draw_networkx_nodes(
        graph,
        pos,
        ax=ax,
        node_color=node_colors,
        node_size=node_size,
        edgecolors=node_edge_color,
        linewidths=node_edge_width,
    )

    nx.draw_networkx_edges(
        graph,
        pos,
        ax=ax,
        edge_color=edge_color,
    )

    if labels is not None:
        if labels == "symbol":
            label_dict = {n: graph.nodes[n]["symbol"] for n in graph.nodes()}
        else:
            label_dict = {n: str(n) for n in graph.nodes()}

        font_colors = {n: _label_color(graph.nodes[n]["color"]) for n in graph.nodes()}

        # Group nodes by font color to minimize draw calls.
        by_color: dict[str, dict] = {}
        for n, lbl in label_dict.items():
            fc = font_colors[n]
            by_color.setdefault(fc, {})[n] = lbl

        for fc, lbl_subset in by_color.items():
            nx.draw_networkx_labels(
                graph,
                pos,
                labels=lbl_subset,
                ax=ax,
                font_color=fc,
            )

    ax.set_axis_off()
    return fig, ax
