"""Dynamic visualization of Theano graphs.

Author: Christof Angermueller <cangermueller@gmail.com>
"""

import os
import shutil
import re
from six import iteritems

from theano.d3viz.formatting import PyDotFormatter

__path__ = os.path.dirname(os.path.realpath(__file__))


def replace_patterns(x, replace):
    """Replace `replace` in string `x`.

    Parameters
    ----------
    s: str
        String on which function is applied
    replace: dict
        `key`, `value` pairs where key is a regular expression and `value` a
        string by which `key` is replaced
    """
    for from_, to in iteritems(replace):
        x = x.replace(str(from_), str(to))
    return x


def escape_quotes(s):
    """Escape quotes in string.

    Parameters
    ----------
    s: str
        String on which function is applied
    """
    s = re.sub(r'''(['"])''', r'\\\1', s)
    return s


def d3viz(fct, outfile, copy_deps=True, *args, **kwargs):
    """Create HTML file with dynamic visualizing of a Theano function graph.

    In the HTML file, the whole graph or single nodes can be moved by drag and
    drop. Zooming is possible via the mouse wheel. Detailed information about
    nodes and edges are displayed via mouse-over events. Node labels can be
    edited by selecting Edit from the context menu.

    Input nodes are colored in green, output nodes in blue. Apply nodes are
    ellipses, and colored depending on the type of operation they perform. Red
    ellipses are transfers from/to the GPU (ops with names GpuFromHost,
    HostFromGpu).

    Edges are black by default. If a node returns a view of an
    input, the input edge will be blue. If it returns a destroyed input, the
    edge will be red.

    Parameters
    ----------
    fct : theano.compile.function_module.Function
        A compiled Theano function, variable, apply or a list of variables.
    outfile : str
        Path to output HTML file.
    copy_deps : bool, optional
        Copy javascript and CSS dependencies to output directory.
    *args : tuple, optional
        Arguments passed to PyDotFormatter.
    *kwargs : dict, optional
        Arguments passed to PyDotFormatter.
    """

    # Create DOT graph
    formatter = PyDotFormatter(*args, **kwargs)
    graph = formatter(fct)
    dot_graph = escape_quotes(graph.create_dot()).replace('\n', '').replace('\r', '')

    # Create output directory if not existing
    outdir = os.path.dirname(outfile)
    if not os.path.exists(outdir):
        os.makedirs(outdir)

    # Read template HTML file
    template_file = os.path.join(__path__, 'html', 'template.html')
    f = open(template_file)
    template = f.read()
    f.close()

    # Copy dependencies to output directory
    src_deps = __path__
    if copy_deps:
        dst_deps = 'd3viz'
        for d in ['js', 'css']:
            dep = os.path.join(outdir, dst_deps, d)
            if not os.path.exists(dep):
                shutil.copytree(os.path.join(src_deps, d), dep)
    else:
        dst_deps = src_deps

    # Replace patterns in template
    replace = {
        '%% JS_DIR %%': os.path.join(dst_deps, 'js'),
        '%% CSS_DIR %%': os.path.join(dst_deps, 'css'),
        '%% DOT_GRAPH %%': dot_graph,
    }
    html = replace_patterns(template, replace)

    # Write HTML file
    with open(outfile, 'w') as f:
        f.write(html)


def d3write(fct, path, *args, **kwargs):
    """Convert Theano graph to pydot graph and write to dot file.

    Parameters
    ----------
    fct : theano.compile.function_module.Function
        A compiled Theano function, variable, apply or a list of variables.
    path: str
        Path to output file
    *args : tuple, optional
        Arguments passed to PyDotFormatter.
    *kwargs : dict, optional
        Arguments passed to PyDotFormatter.
    """

    formatter = PyDotFormatter(*args, **kwargs)
    graph = formatter(fct)
    graph.write_dot(path)
