"""
Provides general visualization utilities. This is NOT related to `rendering`.
"""

from tempfile import NamedTemporaryFile

from pydrake.common import temp_directory

# TODO(eric.cousineau): Move `plot_graphviz` to something more accessible to
# `call_python_client`.

# Use a global variable here because some calls to IPython will actually cause
# an interpreter to be created.  This file needs to be imported BEFORE that
# happens.
try:
    # IPython module may not be available; in which case, we're definitely
    # *not* running as a notebook.
    from IPython import get_ipython
    from IPython.display import SVG, display

    running_as_notebook = get_ipython() and hasattr(get_ipython(), "kernel")
except ModuleNotFoundError:
    running_as_notebook = False


def _plt():
    """Deferred import of plt, so pydrake.all importing doesn't bomb out."""
    # On Ubuntu the Debian package python3-tk is a recommended (but not
    # required) dependency of python3-matplotlib; help users understand
    # that by providing a nicer message upon a failure to import.
    result = None
    try:
        import matplotlib.pyplot as __plt

        result = __plt
    except ImportError as e:
        if e.name == "tkinter":
            result = None
        else:
            raise
    if result is None:
        raise NotImplementedError(
            "On Ubuntu when using the default pyplot configuration (i.e.,"
            " the TkAgg backend) you must 'sudo apt install python3-tk' to"
            " obtain Tk support. Alternatively, you may set MPLBACKEND to"
            " something else (e.g., Qt5Agg)."
        )
    return result


def _pydot():
    """Deferred import of pyplot, so pydrake.all importing doesn't bomb out."""
    import pydot as __pydot

    return __pydot


def plot_graphviz(dot_text):
    """Renders a DOT graph in matplotlib."""
    plt = _plt()
    pydot = _pydot()

    # @ref https://stackoverflow.com/a/18522941/7829525
    # Tried (reason ignored): pydotplus (`pydot` works), networkx
    # (`read_dot` does not work robustly?), pygraphviz (coupled with
    # `networkx`).
    g = pydot.graph_from_dot_data(dot_text)
    if isinstance(g, list):
        # Per Ioannis's follow-up comment in the above link, in pydot >= 1.2.3
        # `graph_from_dot_data` returns a list of graphs.
        # Handle this case for now.
        assert len(g) == 1
        g = g[0]
    if running_as_notebook:
        return display(SVG(g.create_svg()))
    else:
        f = NamedTemporaryFile(suffix=".png", dir=temp_directory())
        g.write_png(f.name)
        plt.axis("off")

        return plt.imshow(plt.imread(f.name), aspect="equal")


def plot_system_graphviz(system, **kwargs):
    """Renders a System's Graphviz representation in `matplotlib`."""
    return plot_graphviz(system.GetGraphvizString(**kwargs))
