"""
Utilities for drawing molecules and cgsmiles graphs.
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import networkx as nx
from .graph_layout import LAYOUT_METHODS
from .drawing_utils import make_graph_edges, make_mapped_edges, make_node_pies

# loosely follow Avogadro / VMD color scheme
ELE_TO_COLOR = {"F": "lightblue",
                "P": "tab:orange",
                "H": "gray",
                "Cl": "tab:green",
                "O": "tab:red",
                "C": "cyan",
                "Na": "pink",
                "N": "blue",
                "Br": "darkred",
                "I": "purple",
                "S": "yellow",
                "Mg": "lightgreen"}

# this list sets the colors for CG fragments. It will be cycled over.
FRAGID_COLORS = ["tab:blue",
                 "tab:red",
                 "tab:orange",
                 "tab:pink",
                 "tab:purple",
                 "tab:cyan",
                 "tab:green",
                 "tab:olive",
                 "tab:brown",
                 "tab:gray"]

DEFAULT_COLOR = "orchid"
DEFAULT_SHARED_COLOR = "gainsboro"

def draw_molecule(graph,
                  ax=None,
                  layout_method='vespr_refined',
                  pos=None,
                  cg_mapping=True,
                  colors=None,
                  labels=None,
                  scale=2,
                  outline=False,
                  use_weights=False,
                  align_with='diag',
                  fontsize=10,
                  text_color='black',
                  edge_widths=3,
                  mapped_edge_width=20,
                  default_bond=1,
                  layout_kwargs={}):
    """
    Draw the graph of a molecule optionally with a coarse-grained
    projection if `cg_mapping` is set to True. The membership of
    atoms to the CG projection is taken from the 'fragid' node attribute.

    Positions or one of three layout methods must be specified. The
    layout options are vespr_refined, vespr, and circular. Note
    that the vespr_refined is slow but yields the best quality results.
    For a quick look vespr is recommended. The drawing function also
    accepts a layout_kwarg dictionary specifiying options to be given
    to the chosen layout method. See also `cgsmiles.graph_layout`.

    For example to draw Benzene using the Martini 3 mapping:

    >>> import matplotlib.pyplot as plt
    >>> from cgsmiles import MoleculeResolver
    >>> from cgsmiles import draw_molecule

    >>> cgsmiles_str = "{[#TC5]1[#TC5][#TC5]1}.{#TC5=[$]cc[$]}"
    >>> resolver = MoleculeResolver.from_string(cgsmiles_str)
    >>> cg_mol, aa_mol = resolver.resolve()
    >>> draw_molecule(aa_mol, layout_method="vespr_refined")
    >>> plt.show()

    The drawing is always of fixed size based on the canvas size. This
    means that molecules drawn on a canvas with the same size will have
    the same dimensions (i.e. bond length, atom size). Using the scale
    argument a drawing can be scaled to make it larger or smaller on
    a given canvas. That means if your molecule is too small or large
    redraw it setting a different scale parameter.

    Parameters
    ----------
    ax: matplotlib.axes.Axes
        mpl axis object
    layout_method: str
        choice of vespr, vespr_refined, circular
        (default: 'vespr_refined')
    pos: dict
        a dict mapping nodes to 2D positions
    cg_mapping: bool
        draw outline for the CG mapping (default: True)
    colors: dict
        a dict mapping nodes to colors or fragids to colors
        depending on if cg_mapping is True
    labels: list
        list of node_labels; must follow the order of nodes in graph
    scale: float
        scale the drawing relative to the total canvas size
    outline: bool
        draw an outline around each node
    use_weights: bool
        color nodes according to weight attribute (default: False)
    align_with: str or :class:`numpy.ndarray`
        align the longest distance in molecule with one of x, y, diag
        or a custom axis as numpy 2D array
    fontsize: float
        fontsize of labels
    text_color: str
        color of the node labels (default: black)
    edge_widths: float
        the width of the bonds
    mapped_edge_widths: float
        the width of the mapped projection
    default_bond: float
        default bond length (default: 1)
    layout_kwargs: dict
        dict with arguments passed to the layout methods

    Returns
    -------
    matplotlib.axes.Axes
        the updated axis object
    dict
        a dict of positions
    """
    # check input args
    if pos is None and layout_method is None:
        msg = "You need to provide either positions or a layout method."
        raise ValueError(msg)

    # scaling cannot be negative
    if scale < 0:
        raise ValueError('scale should not be negative')

    # generate figure in case we don't get one
    if not ax:
        ax = plt.gca()

    # scale edge widths and fontsize
    if scale:
        edge_widths = edge_widths * scale
        mapped_edge_width = mapped_edge_width * scale
        fontsize = fontsize * scale

    # default labels are the element names
    if labels is None:
        labels = nx.get_node_attributes(graph, 'element')

    # collect the fragids if CG projection is to be drawn
    if cg_mapping:
        ids = nx.get_node_attributes(graph, 'fragid')
        id_set = set()
        for fragid in ids.values():
            id_set |= set(fragid)

    # assign color defaults
    if colors is None and cg_mapping:
        colors = {fragid: FRAGID_COLORS[fragid % len(FRAGID_COLORS)]
                  for fragid in id_set}
    elif colors is None:
        colors = {node: ELE_TO_COLOR.get(ele, DEFAULT_COLOR)
                  for node, ele in nx.get_node_attributes(graph, 'element').items()}

    # if no positions are given generate layout
    bbox = ax.get_position(True)

    # some axis magic
    fig_width_inch, fig_height_inch = ax.figure.get_size_inches()
    w = bbox.width*fig_width_inch/scale
    h = bbox.height*fig_height_inch/scale

    # compute angle for axis alignment
    _keyword_to_axis = {'diag': np.array([w, h]),
                        'x': np.array([1, 0]),
                        'y': np.array([0, 1])}
    if isinstance(align_with, str):
        align_with = _keyword_to_axis[align_with]

    # generate inital positions
    if not pos:
        pos = LAYOUT_METHODS[layout_method](graph,
                                            default_bond=default_bond,
                                            align_with=align_with,
                                            **layout_kwargs)



    # generate starting and stop positions for edges
    edges, arom_edges, plain_edges = make_graph_edges(graph, pos)

    # draw the edges
    ax.add_collection(LineCollection(edges,
                                     color='black',
                                     linewidths=edge_widths,
                                     zorder=2))
    ax.add_collection(LineCollection(arom_edges,
                                     color='black',
                                     linestyle='dotted',
                                     linewidths=edge_widths, zorder=2))

    # generate the edges from the mapping
    if cg_mapping:
        mapped_edges = make_mapped_edges(graph, plain_edges)
        for fragids, frag_edges in mapped_edges.items():
            if len(fragids) == 1:
                color = colors.get(*fragids)
            else:
                color = DEFAULT_SHARED_COLOR
            ax.add_collection(LineCollection(frag_edges,
                                             color=color,
                                             linewidths=mapped_edge_width,
                                             zorder=1,
                                             alpha=0.5))

    # now we draw nodes
    for slices, pie_kwargs in make_node_pies(graph,
                                             pos,
                                             cg_mapping,
                                             colors,
                                             outline=outline,
                                             radius=default_bond/3.,
                                             use_weights=use_weights,
                                             linewidth=edge_widths):
        p, _ = ax.pie(slices, **pie_kwargs)
        for pie in p:
            pie.set_zorder(3)

    # add node texts
    zorder=4
    for idx, label in labels.items():
        x, y = pos[idx]
        ax.text(x, y, label,
                zorder=zorder,
                fontsize=fontsize,
                verticalalignment='center_baseline',
                horizontalalignment='center',
                color=text_color)
        zorder+=1

    # compute initial view
    w = bbox.width
    h = bbox.height

    fig_width_inch, fig_height_inch = ax.figure.get_size_inches()
    w = bbox.width*fig_width_inch/scale
    h = bbox.height*fig_height_inch/scale
    ax.set_xlim(-w, w)
    ax.set_ylim(-h, h)

    return ax, pos
