from typing import Union

import networkx as nx
import matplotlib.pyplot as plt
from networkx.drawing.nx_pydot import graphviz_layout

from .net import EHMNet, EHM2Net, EHM2Tree
from .utils import to_nx_graph


def plot_net(net: Union[EHMNet, EHM2Net], ax: plt.Axes = None, annotate=True):
    """Plot the net.

    Parameters
    ----------
    net : :class:`~.EHMNet` | :class:`~.EHM2Net`
        The net to plot.
    ax : :class:`matplotlib.axes.Axes`
        Axes on which to plot the net. If ``None``, a new figure and axes will be created.
    annotate : :class:`bool`
        Flag that dictates whether to draw node and edge labels on the plotted net. The default is ``True``
    """
    if ax is None:
        fig = plt.figure()
        ax = fig.gca()
    g = to_nx_graph(net)
    pos = graphviz_layout(g, prog="dot")
    nx.draw(g, pos, ax=ax, node_size=0)
    if annotate:
        labels = dict()
        for n in g.nodes:
            t = g.nodes[n]['track']
            s = str(g.nodes[n]['identity']) if len(g.nodes[n]['identity']) else 'Ø'
            if t is not None:
                labels[n] = '{{{}, {}}}'.format(t, s)
            else:
                labels[n] = 'Ø'
        pos_labels = {}
        for node, coords in pos.items():
            pos_labels[node] = (coords[0] + 10, coords[1])
        nx.draw_networkx_labels(g, pos_labels, ax=ax, labels=labels, horizontalalignment='left')
        edge_labels = nx.get_edge_attributes(g, 'detections')
        for key in edge_labels:
            edge_labels[key] = str(edge_labels[key]).replace('{', '').replace('}', '')
        nx.draw_networkx_edge_labels(g, pos, edge_labels=edge_labels)


def plot_tree(tree: EHM2Tree, ax: plt.Axes = None, annotate=True):
    """Plot the tree.

    Parameters
    ----------
    tree : :class:`~.EHM2Tree`
        The tree to plot.
    ax: :class:`matplotlib.axes.Axes`
        Axes on which to plot the tree. If ``None``, a new figure and axes will be created.
    annotate : :class:`bool`
        Flag that dictates whether to draw node labels on the plotted tree. The default is ``True``
    """
    if ax is None:
        fig = plt.figure()
        ax = fig.gca()
    g = to_nx_graph(tree)
    pos = graphviz_layout(g, prog="dot")
    nx.draw(g, pos, ax=ax)
    labels = {n: g.nodes[n]['track'] for n in g.nodes}  # if g.nodes[n]['leaf']}
    pos_labels = {}
    if annotate:
        for node, coords in pos.items():
            # if g.nodes[node]['leaf']:
            pos_labels[node] = (coords[0], coords[1])
        nx.draw_networkx_labels(g, pos_labels, ax=ax, labels=labels, font_color='white')
