# AUTOGENERATED! DO NOT EDIT! File to edit: display.ipynb (unless otherwise specified).

__all__ = ['show', 'pretty']

# Cell
import gtsam
import graphviz
from .discrete import Variables

# Cell
class show(graphviz.Source):
    """ Display an object with a dot method as a graph."""

    def __init__(self, obj, *args, **kwargs):
        """Construct from object with 'dot' method."""
        # This small class takes an object, calls its dot function, and uses the
        # resulting string to initialize a graphviz.Source instance. This in turn
        # has a _repr_mimebundle_ method, which then renders it in the notebook.
        if args and isinstance(args[0], Variables):
            assert len(args) == 1, "Variables must be only argument."
            keyFormatter = args[0].keyFormatter()
            super().__init__(obj.dot(keyFormatter))
        else:
            super().__init__(obj.dot(*args, **kwargs))

# Cell
class pretty:
    """Render an object as markdown with optional arguments."""

    def __init__(self, obj, *args):
        if args and isinstance(args[0], Variables):
            assert len(args) == 1, "Variables must be only argument."
            variables = args[0]
            if isinstance(obj, gtsam.DiscreteValues):
                self._md = variables.values_markdown(obj)
            else:
                self._md = obj._repr_markdown_(
                    variables.keyFormatter(), variables.names())
        else:
            if isinstance(obj, gtsam.DiscreteValues):
                self._md = f"{obj}"
            else:
                self._md = obj._repr_markdown_(*args)

    def _repr_markdown_(self):
        return self._md
