"""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.5,
):
    """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,
        width=1.5,
    )

    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


def draw_matches(
    graph: nx.Graph | str,
    matches: list,
    *,
    highlight_color: object = (0.6, 1.0, 0.6),
    highlight_node_scale: float = 2.2,
    highlight_edge_width: float = 8.0,
    ax=None,
    pos: dict | None = 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.5,
):
    """Draw a COGILES graph with pattern matches highlighted.

    Matched nodes get a larger colored circle behind them (halo effect)
    and matched edges get a fatter colored line behind them.

    When *ax* is ``None``, the full graph is drawn with highlights.
    When *ax* is provided (from a previous call), only highlights are
    added — this allows chaining multiple calls with different colors.

    Parameters
    ----------
    graph : nx.Graph or str
        A NetworkX graph with COGILES node attributes, or a COGILES string.
    matches : list[Match]
        Pattern matches to highlight (from :func:`find_pattern`).
    highlight_color : color, optional
        Color for the highlight halos and edges (default: light green).
        Accepts any matplotlib color specification.
    highlight_node_scale : float, optional
        Size multiplier for highlight circles relative to *node_size*
        (default 1.8).
    highlight_edge_width : float, optional
        Line width of the highlight edges (default 6.0).
    ax : matplotlib.axes.Axes, optional
        Axes to draw on.  If *None*, a new figure is created and the
        full graph is drawn.  If provided, only highlights are added.
    pos : dict, optional
        Precomputed node positions ``{node: (x, y)}``.  Required for
        consistent positioning when chaining calls with non-deterministic
        layouts (e.g. ``"spring"``).  If *None*, positions are computed
        from *layout*.
    layout : str, optional
        Layout algorithm name (default ``"kamada_kawai"``).
    labels : str or None, optional
        Node label style (default ``None``).  Only used when *ax* is
        ``None`` (i.e. when drawing the full graph).
    node_size : int, optional
        Size of the drawn nodes (default 300).
    edge_color : str, optional
        Color of the regular 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.
    """
    from matplotlib.collections import LineCollection

    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 pos is None:
        pos = LAYOUT_MAP[layout](graph)

    # -- Collect matched nodes and edges --

    matched_nodes: set[int] = set()
    matched_edges: set[tuple[int, int]] = set()
    for m in matches:
        matched_nodes |= m.nodes
        for n1 in m.nodes:
            for n2 in m.nodes:
                if n1 < n2 and graph.has_edge(n1, n2):
                    matched_edges.add((n1, n2))

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

    # -- Draw highlight layers (low z-order → behind graph) --

    if matched_edges:
        segments = [
            [(pos[u][0], pos[u][1]), (pos[v][0], pos[v][1])]
            for u, v in matched_edges
        ]
        lc = LineCollection(
            segments,
            colors=[highlight_color] * len(segments),
            linewidths=highlight_edge_width,
            zorder=-9,
        )
        ax.add_collection(lc)

    if matched_nodes:
        x = [pos[n][0] for n in matched_nodes]
        y = [pos[n][1] for n in matched_nodes]
        ax.scatter(
            x, y,
            s=node_size * highlight_node_scale,
            c=[highlight_color] * len(matched_nodes),
            edgecolors="none",
            zorder=-10,
        )

    # -- Draw the regular graph on top (only on first call) --

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

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

        edge_col = nx.draw_networkx_edges(
            graph,
            pos,
            ax=ax,
            edge_color=edge_color,
            width=1.5,
        )
        if edge_col is not None:
            edge_col.set_zorder(2)

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

            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()
            }

            by_fc: dict[str, dict] = {}
            for n, lbl in label_dict.items():
                fc = font_colors[n]
                by_fc.setdefault(fc, {})[n] = lbl

            for fc, lbl_subset in by_fc.items():
                texts = nx.draw_networkx_labels(
                    graph,
                    pos,
                    labels=lbl_subset,
                    ax=ax,
                    font_color=fc,
                )
                for t in texts.values():
                    t.set_zorder(4)

        ax.set_axis_off()

    return fig, ax
