import networkx as nx
import re
import functools
import io
import warnings
import inspect
from .table import Table
from .dependencies import unite_master_parts

try:
    from matplotlib import pyplot as plt
    plot_active = True
except:
    plot_active = False

try:
    from networkx.drawing.nx_pydot import pydot_layout
    diagram_active = True
except:
    diagram_active = False

from .user_tables import Manual, Imported, Computed, Lookup
from .part_table import Part
from .auto_tables import AutoComputed, AutoImported
from .settings_table import Settingstable
from .errors import DataJointError
from .table import lookup_class_name


user_table_classes = (
    Manual, Lookup, Computed, Imported, Part, AutoComputed, AutoImported,
    Settingstable
)


class _AliasNode:
    """
    special class to indicate aliased foreign keys
    """
    pass


def _get_tier(table_name):
    if not table_name.startswith('`'):
        return _AliasNode
    else:
        try:
            return next(tier for tier in user_table_classes
                        if re.fullmatch(tier.tier_regexp, table_name.split('`')[-2]))
        except StopIteration:
            return None


if not diagram_active:
    class Diagram:
        """
        Entity relationship diagram, currently disabled due to the lack of required packages: matplotlib and pygraphviz.

        To enable Diagram feature, please install both matplotlib and pygraphviz. For instructions on how to install
        these two packages, refer to http://docs.datajoint.io/setup/Install-and-connect.html#python and
        http://tutorials.datajoint.io/setting-up/datajoint-python.html
        """

        def __init__(self, *args, **kwargs):
            warnings.warn('Please install matplotlib and pygraphviz libraries to enable the Diagram feature.')

else:
    class Diagram(nx.DiGraph):
        """
        Entity relationship diagram.

        Usage:

        >>>  diag = Diagram(source)

        source can be a base relation object, a base relation class, a schema, or a module that has a schema.

        >>> diag.draw()

        draws the diagram using pyplot

        diag1 + diag2  - combines the two diagrams.
        diag + n   - expands n levels of successors
        diag - n   - expands n levels of predecessors
        Thus dj.Diagram(schema.Table)+1-1 defines the diagram of immediate ancestors and descendants of schema.Table

        Note that diagram + 1 - 1  may differ from diagram - 1 + 1 and so forth.
        Only those tables that are loaded in the connection object are displayed
        """
        def __init__(self, source, context=None):

            if isinstance(source, Diagram):
                # copy constructor
                self.nodes_to_show = set(source.nodes_to_show)
                self.context = source.context
                super().__init__(source)
                return

            # get the caller's context
            if context is None:
                frame = inspect.currentframe().f_back
                self.context = dict(frame.f_globals, **frame.f_locals)
                del frame
            else:
                self.context = context

            # find connection in the source
            try:
                connection = source.connection
            except AttributeError:
                try:
                    connection = source.schema.connection
                except AttributeError:
                    raise DataJointError('Could not find database connection in %s' % repr(source[0]))

            # initialize graph from dependencies
            connection.dependencies.load()
            super().__init__(connection.dependencies)

            # Enumerate nodes from all the items in the list
            self.nodes_to_show = set()
            try:
                self.nodes_to_show.add(source.full_table_name)
            except AttributeError:
                try:
                    database = source.database
                except AttributeError:
                    try:
                        database = source.schema.database
                    except AttributeError:
                        raise DataJointError('Cannot plot Diagram for %s' % repr(source))
                for node in self:
                    if node.startswith('`%s`' % database):
                        self.nodes_to_show.add(node)

        @classmethod
        def from_sequence(cls, sequence):
            """
            The join Diagram for all objects in sequence
            :param sequence: a sequence (e.g. list, tuple)
            :return: Diagram(arg1) + ... + Diagram(argn)
            """
            return functools.reduce(lambda x, y: x+y, map(Diagram, sequence))

        def add_parts(self):
            """
            Adds to the diagram the part tables of tables already included in the diagram
            :return:
            """
            def is_part(part, master):
                """
                :param part:  `database`.`table_name`
                :param master:   `database`.`table_name`
                :return: True if part is part of master.
                """
                part = [s.strip('`') for s in part.split('.')]
                master = [s.strip('`') for s in master.split('.')]
                return master[0] == part[0] and master[1] + '__' == part[1][:len(master[1])+2]

            self = Diagram(self)  # copy
            self.nodes_to_show.update(n for n in self.nodes() if any(is_part(n, m) for m in self.nodes_to_show))
            return self

        def topological_sort(self):
            """ :return:  list of nodes in topological order """
            return unite_master_parts(list(nx.algorithms.dag.topological_sort(
                nx.DiGraph(self).subgraph(self.nodes_to_show))))

        def __add__(self, arg):
            """
            :param arg: either another Diagram or a positive integer.
            :return: Union of the diagrams when arg is another Diagram
                     or an expansion downstream when arg is a positive integer.
            """
            self = Diagram(self)   # copy
            try:
                self.nodes_to_show.update(arg.nodes_to_show)
            except AttributeError:
                try:
                    self.nodes_to_show.add(arg.full_table_name)
                except AttributeError:
                    for i in range(arg):
                        new = nx.algorithms.boundary.node_boundary(self, self.nodes_to_show)
                        if not new:
                            break
                        # add nodes referenced by aliased nodes
                        new.update(nx.algorithms.boundary.node_boundary(self, (a for a in new if a.isdigit())))
                        self.nodes_to_show.update(new)
            return self

        def __sub__(self, arg):
            """
            :param arg: either another Diagram or a positive integer.
            :return: Difference of the diagrams when arg is another Diagram or
                     an expansion upstream when arg is a positive integer.
            """
            self = Diagram(self)   # copy
            try:
                self.nodes_to_show.difference_update(arg.nodes_to_show)
            except AttributeError:
                try:
                    self.nodes_to_show.remove(arg.full_table_name)
                except AttributeError:
                    for i in range(arg):
                        graph = nx.DiGraph(self).reverse()
                        new = nx.algorithms.boundary.node_boundary(graph, self.nodes_to_show)
                        if not new:
                            break
                        # add nodes referenced by aliased nodes
                        new.update(nx.algorithms.boundary.node_boundary(graph, (a for a in new if a.isdigit())))
                        self.nodes_to_show.update(new)
            return self

        def __mul__(self, arg):
            """
            Intersection of two diagrams
            :param arg: another Diagram
            :return: a new Diagram comprising nodes that are present in both operands.
            """
            self = Diagram(self)   # copy
            self.nodes_to_show.intersection_update(arg.nodes_to_show)
            return self

        def _make_graph(self):
            """
            Make the self.graph - a graph object ready for drawing
            """
            # mark "distinguished" tables, i.e. those that introduce new primary key attributes
            for name in self.nodes_to_show:
                foreign_attributes = set(
                    attr for p in self.in_edges(name, data=True) for attr in p[2]['attr_map'] if p[2]['primary'])
                self.nodes[name]['distinguished'] = (
                        'primary_key' in self.nodes[name] and foreign_attributes < self.nodes[name]['primary_key'])
            # include aliased nodes that are sandwiched between two displayed nodes
            gaps = set(nx.algorithms.boundary.node_boundary(self, self.nodes_to_show)).intersection(
                nx.algorithms.boundary.node_boundary(nx.DiGraph(self).reverse(), self.nodes_to_show))
            nodes = self.nodes_to_show.union(a for a in gaps if a.isdigit)
            # construct subgraph and rename nodes to class names
            graph = nx.DiGraph(nx.DiGraph(self).subgraph(nodes))
            nx.set_node_attributes(graph, name='node_type', values={n: _get_tier(n) for n in graph})
            # relabel nodes to class names
            mapping = {node: lookup_class_name(node, self.context) or node for node in graph.nodes()}
            new_names = [mapping.values()]
            if len(new_names) > len(set(new_names)):
                raise DataJointError('Some classes have identical names. The Diagram cannot be plotted.')
            nx.relabel_nodes(graph, mapping, copy=False)
            return graph

        def make_dot(self):

            graph = self._make_graph()
            graph.nodes()

            scale = 1.2   # scaling factor for fonts and boxes
            label_props = {  # http://matplotlib.org/examples/color/named_colors.html
                None: dict(
                    shape='circle', color="#FFFF0040",
                    fontcolor='yellow', fontsize=round(scale*8),
                    size=0.4*scale, fixed=False),
                _AliasNode: dict(
                    shape='circle', color="#FF880080",
                    fontcolor='#FF880080', fontsize=round(scale*0),
                    size=0.05*scale, fixed=True),
                Manual: dict(
                    shape='box', color='greenyellow',
                    fontcolor='black', fontsize=round(scale*10),
                    size=0.4*scale, fixed=False),
                Settingstable: dict(
                    shape='box', color='coral',
                    fontcolor='black', fontsize=round(scale*8),
                    size=0.4*scale, fixed=False),
                Lookup: dict(
                    shape='plaintext', color='azure4',
                    fontcolor='black', fontsize=round(scale*8),
                    size=0.4*scale, fixed=False),
                Computed: dict(
                    shape='ellipse', color='magenta',
                    fontcolor='black', fontsize=round(scale*10),
                    size=0.3*scale, fixed=True),
                Imported: dict(
                    shape='ellipse', color='skyblue',
                    fontcolor='black', fontsize=round(scale*10),
                    size=0.4*scale, fixed=False),
                AutoComputed: dict(
                    shape='ellipse', color='magenta1',
                    fontcolor='black', fontsize=round(scale*10),
                    size=0.3*scale, fixed=True),
                AutoImported: dict(
                    shape='ellipse', color='skyblue1',
                    fontcolor='black', fontsize=round(scale*10),
                    size=0.4*scale, fixed=False),
                Part: dict(
                    shape='plaintext', color='azure4',
                    fontcolor='black', fontsize=round(scale*8),
                    size=0.1*scale, fixed=False)}
            node_props = {node: label_props[d['node_type']] for node, d in dict(graph.nodes(data=True)).items()}

            dot = nx.drawing.nx_pydot.to_pydot(graph)
            for node in dot.get_nodes():
                node.set_shape('circle')
                name = node.get_name().strip('"')
                props = node_props[name]
                node.set_fontsize(props['fontsize'])
                node.set_fontcolor(props['fontcolor'])
                node.set_shape(props['shape'])
                node.set_fontname('arial')
                node.set_fixedsize('shape' if props['fixed'] else False)
                node.set_width(props['size'])
                node.set_height(props['size'])
                if name.split('.')[0] in self.context:
                    cls = eval(name, self.context)
                    assert(issubclass(cls, Table))
                    description = cls().describe(context=self.context, printout=False).split('\n')
                    description = (
                        '-'*30 if q.startswith('---') else q.replace('->', '&#8594;') if '->' in q else q.split(':')[0]
                        for q in description if not q.startswith('#'))
                    node.set_tooltip('&#13;'.join(description))
                node.set_label("<<u>"+name+"</u>>" if node.get('distinguished') == 'True' else name)
                node.set_color(props['color'])
                node.set_style('filled')

            for edge in dot.get_edges():
                # see http://www.graphviz.org/content/attrs
                src = edge.get_source().strip('"')
                dest = edge.get_destination().strip('"')
                props = graph.get_edge_data(src, dest)
                edge.set_color('#00000040')
                edge.set_style('solid' if props['primary'] else 'dashed')
                master_part = graph.nodes[dest]['node_type'] is Part and dest.startswith(src+'.')
                edge.set_weight(3 if master_part else 1)
                edge.set_arrowhead('none')
                edge.set_penwidth(.75 if props['multi'] else 2)

            return dot

        def make_svg(self):
            from IPython.display import SVG
            return SVG(self.make_dot().create_svg())

        def make_png(self):
            return io.BytesIO(self.make_dot().create_png())

        def make_image(self):
            if plot_active:
                return plt.imread(self.make_png())
            else:
                raise DataJointError("pyplot was not imported")

        def _repr_svg_(self):
            return self.make_svg()._repr_svg_()

        def draw(self):
            if plot_active:
                plt.imshow(self.make_image())
                plt.gca().axis('off')
                plt.show()
            else:
                raise DataJointError("pyplot was not imported")

        def save(self, filename, format=None):
            if format is None:
                if filename.lower().endswith('.png'):
                    format = 'png'
                elif filename.lower().endswith('.svg'):
                    format = 'svg'
            if format.lower() == 'png':
                with open(filename, 'wb') as f:
                    f.write(self.make_png().getbuffer().tobytes())
            elif format.lower() == 'svg':
                with open(filename, 'w') as f:
                    f.write(self.make_svg().data)
            else:
                raise DataJointError('Unsupported file format')

        @staticmethod
        def _layout(graph, **kwargs):
            return pydot_layout(graph, prog='dot', **kwargs)
