"""Graphing utilities for EventStreams"""
from __future__ import absolute_import, division, print_function

from functools import partial
import os
import re


def _clean_text(text, match=None):
    ''' Clean text, remove forbidden characters.
    '''
    # all non alpha numeric characters, except for _ and :
    # replace them with space
    # the + condenses a group of consecutive characters all into one space
    # (rather than assigning a space to each)
    if match is None:
        match = '[^a-zA-Z0-9_:]+'
    text = re.sub(match, ' ', text)
    # now replace the colon with semicolon
    text = re.sub(":", ";", text)
    return text


def build_node_set(node, s=None):
    """Build a set of all the nodes in a streamz graph

    Parameters
    ----------
    node : Stream
        The node to use as a starting point for building the set
    s : set or None
        The set to put the nodes into. If None return a new set full of nodes

    Returns
    -------
    s : set
        The set of nodes in the graph

    """
    if s is None:
        s = set()
    if node is None or (
        node in s
        and all(n in s for n in node.upstreams)
        and all(n in s for n in node.downstreams)
    ):
        return
    new_nodes = {n for n in node.downstreams}
    new_nodes.update(node.upstreams)
    new_nodes.add(node)
    s.update(new_nodes)
    [build_node_set(n, s) for n in list(new_nodes)]
    return s


def create_graph(node, graph):
    """Create networkx graph of the pipeline

    Parameters
    ----------
    node : Stream
        The node to start from
    graph : networkx.DiGraph
        The graph to fill with nodes

    Returns
    -------

    """
    # Step 1 build a set of all the nodes
    node_set = build_node_set(node)

    # Step 2 for each node in the set add to the graph
    for n in node_set:
        t = hash(n)
        graph.add_node(
            t,
            label=_clean_text(str(n)),
            shape=n._graphviz_shape,
            orientation=str(n._graphviz_orientation),
            style=n._graphviz_style,
            fillcolor=n._graphviz_fillcolor,
        )

    # Step 3 for each node establish its edges
    for n in node_set:
        t = hash(n)
        for nn in n.upstreams:
            tt = hash(nn)
            graph.add_edge(tt, t)

        downstreams = n.downstreams
        for i, nn in enumerate(downstreams):
            tt = hash(nn)
            if len(downstreams) > 1:
                graph.add_edge(t, tt, label=str(i))
            else:
                graph.add_edge(t, tt)

    # Step 4 destroy set
    del node_set


def readable_graph(graph):
    """Create human readable version of this object's task graph.

    Parameters
    ----------
    graph: nx.DiGraph instance
        The networkx graph representing the pipeline
    """
    import networkx as nx

    mapping = {k: "{}".format(graph.nodes[k]["label"]) for k in graph}
    idx_mapping = {}
    for k, v in mapping.items():
        if v in idx_mapping.keys():
            idx_mapping[v] += 1
            mapping[k] += "-{}".format(idx_mapping[v])
        else:
            idx_mapping[v] = 0

    gg = {k: v for k, v in mapping.items()}
    rg = nx.relabel_nodes(graph, gg, copy=True)
    return rg


def to_graphviz(graph, **graph_attr):
    import graphviz

    digraph_kwargs = {'name', 'comment', 'filename',
                      'format', 'engine', 'encoding',
                      'graph_attr', 'node_attr', 'edge_attr',
                      'body', 'strict', 'directory'}
    if not digraph_kwargs.intersection(graph_attr):
        graph_attr = dict(graph_attr=graph_attr)

    gvz = graphviz.Digraph(**graph_attr)
    for node, attrs in graph.nodes.items():
        gvz.node(node, **attrs)
    for edge, attrs in graph.edges().items():
        gvz.edge(edge[0], edge[1], **attrs)
    return gvz


def visualize(node, filename="mystream.png", **kwargs):
    """
    Render a task graph using dot.

    If `filename` is not None, write a file to disk with that name in the
    format specified by `format`.  `filename` should not include an extension.

    Parameters
    ----------
    node : Stream instance
        The stream to display.
    filename : str or None, optional
        The name (without an extension) of the file to write to disk.  If
        `filename` is None, no file will be written, and we communicate with
        dot using only pipes.  Default is 'mydask'.
    format : {'png', 'pdf', 'dot', 'svg', 'jpeg', 'jpg'}, optional
        Format in which to write output file.  Default is 'png'.

    Returns
    -------
    result : None or IPython.display.Image or IPython.display.SVG  (See below.)

    Notes
    -----
    If IPython is installed, we return an IPython.display object in the
    requested format.  If IPython is not installed, we just return None.

    We always return None if format is 'pdf' or 'dot', because IPython can't
    display these formats natively. Passing these formats with filename=None
    will not produce any useful output.

    See Also
    --------
    streams.graph.readable_graph
    """
    import networkx as nx

    nx_g = nx.DiGraph()
    create_graph(node, nx_g)
    rg = readable_graph(nx_g)
    g = to_graphviz(rg, **kwargs)

    fmts = [".png", ".pdf", ".dot", ".svg", ".jpeg", ".jpg"]
    if filename is None:
        format = "png"

    elif any(filename.lower().endswith(fmt) for fmt in fmts):
        filename, format = os.path.splitext(filename)
        format = format[1:].lower()

    else:
        format = "png"

    data = g.pipe(format=format)
    if not data:
        raise RuntimeError(
            "Graphviz failed to properly produce an image. "
            "This probably means your installation of graphviz "
            "is missing png support. See: "
            "https://github.com/ContinuumIO/anaconda-issues/"
            "issues/485 for more information."
        )

    display_cls = _get_display_cls(format)

    if not filename:
        return display_cls(data=data)

    full_filename = ".".join([filename, format])
    with open(full_filename, "wb") as f:
        f.write(data)

    return display_cls(filename=full_filename)


IPYTHON_IMAGE_FORMATS = frozenset(["jpeg", "png"])
IPYTHON_NO_DISPLAY_FORMATS = frozenset(["dot", "pdf"])


def _get_display_cls(format):
    """
    Get the appropriate IPython display class for `format`.

    Returns `IPython.display.SVG` if format=='svg', otherwise
    `IPython.display.Image`.

    If IPython is not importable, return dummy function that swallows its
    arguments and returns None.
    """
    dummy = lambda *args, **kwargs: None
    try:
        import IPython.display as display
    except ImportError:
        # Can't return a display object if no IPython.
        return dummy

    if format in IPYTHON_NO_DISPLAY_FORMATS:
        # IPython can't display this format natively, so just return None.
        return dummy
    elif format in IPYTHON_IMAGE_FORMATS:
        # Partially apply `format` so that `Image` and `SVG` supply a uniform
        # interface to the caller.
        return partial(display.Image, format=format)
    elif format == "svg":
        return display.SVG
    else:
        raise ValueError("Unknown format '%s' passed to `dot_graph`" % format)