from typing import (
    Any,
    Optional,
    Sequence,
)
from collections.abc import Hashable
import copy
from contextlib import contextmanager
import numpy as np
import pandas as pd


style_leaves = (
    "cmap",
    "color",
    "size",
    "edgecolor",
    "facecolor",
    "linewidth",
    "linestyle",
    "alpha",
    "zorder",
    "tension",
    "looptension",
    "loopmaxangle",
    "paralleloffset",
    "offset",
    "rotate",
    "marker",
    "waypoints",
    "horizontalalignment",
    "verticalalignment",
    "boxstyle",
    "hpadding",
    "vpadding",
    "hmargin",
    "vmargin",
    "ports",
    "extend",
)

# These properties are not allowed to be rotated (global throughout the graph).
# This might change in the future as the API improves.
nonrotating_leaves = (
    "paralleloffset",
    "looptension",
    "loopmaxangle",
    "vertexpadding",
    "extend",
)


default = {
    "vertex": {
        "size": 20,
        "facecolor": "black",
        "marker": "o",
        "label": {
            "color": "white",
            "horizontalalignment": "center",
            "verticalalignment": "center",
            "hpadding": 18,
            "vpadding": 12,
        },
    },
    "edge": {
        "linewidth": 1.5,
        "linestyle": "-",
        "color": "black",
        "curved": False,
        "paralleloffset": 3,
        "tension": 1,
        "looptension": 4,
        "loopmaxangle": 60,
        "label": {
            "horizontalalignment": "center",
            "verticalalignment": "center",
            "rotate": False,
            "bbox": {
                "boxstyle": "round",
                "facecolor": "white",
                "edgecolor": "none",
            },
        },
        "arrow": {
            "marker": "|>",
            "width": 8,
        },
    },
    "grouping": {
        "facecolor": ["grey", "steelblue", "tomato"],
        "edgecolor": "black",
        "linewidth": 1.5,
        "alpha": 0.5,
        "vertexpadding": 18,
    },
}


def copy_with_deep_values(style):
    """Make a deep copy of the style dict but do not create copies of the keys."""
    newdict = {}
    for key, value in style.items():
        if isinstance(value, dict):
            newdict[key] = copy_with_deep_values(value)
        else:
            newdict[key] = copy.copy(value)
    return newdict


hollow = copy_with_deep_values(default)
hollow["vertex"]["color"] = None
hollow["vertex"]["facecolor"] = "none"
hollow["vertex"]["edgecolor"] = "black"
hollow["vertex"]["linewidth"] = 1.5
hollow["vertex"]["marker"] = "r"
hollow["vertex"]["size"] = "label"
hollow["vertex"]["label"]["color"] = "black"

tree = copy_with_deep_values(default)
tree["vertex"]["size"] = 0
tree["vertex"]["alpha"] = 0
tree["edge"]["linewidth"] = 2.5
tree["vertex"]["label"]["bbox"] = {
    "boxstyle": "square,pad=0.5",
    "facecolor": "white",
    "edgecolor": "none",
}
tree["vertex"]["label"]["color"] = "black"
tree["vertex"]["label"]["size"] = 12
tree["vertex"]["label"]["verticalalignment"] = "center"
tree["vertex"]["label"]["hmargin"] = 10


styles = {
    "default": default,
    "hollow": hollow,
    "tree": tree,
}


stylename = "default"


current = copy_with_deep_values(styles["default"])


def get_stylename():
    """Return the name of the current iplotx style."""
    return str(stylename)


def get_style(name: str = "", *args) -> dict[str, Any]:
    """Get a *deep copy* of the chosen style.

    Parameters:
        name: The name of the style to get. If empty, the current style is returned.
            Substyles can be obtained by using a dot notation, e.g. "default.vertex".
            If "name" starts with a dot, it means a substyle of the current style.
        *args: A single argument is accepted. If present, this value (usually a
            dictionary) is returned if the queried style is not found. For example,
            get_style(".nonexistent") raises an Exception but
            get_style("nonexistent", {}) does not, returning an empty dict instead.
    Returns:
        The requected style or substyle.

    NOTE: The deep copy is a little different from standard deep copies. Here, keys
        (which need to be hashables) are never copied, but values are. This can be
        useful for hashables that change hash upon copying, such as Biopython's
        tree nodes.
    """
    if len(args) > 1:
        raise ValueError("get_style() accepts at most one additional argument.")

    namelist = name.split(".")
    style = styles
    for i, namei in enumerate(namelist):
        if (i == 0) and (namei == ""):
            style = current
        else:
            if namei in style:
                style = style[namei]
            # NOTE: if asking for a nonexistent, non-leaf style
            # give the benefit of the doubt and set an empty dict
            # which will not fail unless the uder tries to enter it
            elif namei not in style_leaves:
                style = {}
            elif len(args) > 0:
                return args[0]
            else:
                raise KeyError(f"Style not found: {name}")

    style = copy_with_deep_values(style)
    return style


# The following is inspired by matplotlib's style library
# https://github.com/matplotlib/matplotlib/blob/v3.10.3/lib/matplotlib/style/core.py#L45
def use(style: Optional[str | dict | Sequence] = None, **kwargs):
    """Use iplotx style setting for a style specification.

    The style name of 'default' is reserved for reverting back to
    the default style settings.

    Parameters:
        style: A style specification, currently either a name of an existing style
            or a dict with specific parts of the style to override. The string
            "default" resets the style to the default one. If this is a sequence,
            each style is applied in order.
        **kwargs: Additional style changes to be applied at the end of any style.
    """
    try:
        import networkx as nx
    except ImportError:
        nx = None

    global current

    def _sanitize_leaves(style: dict):
        for key, value in style.items():
            if key in style_leaves:
                if nx is not None:
                    if isinstance(value, nx.classes.reportviews.NodeView):
                        style[key] = dict(value)
                    elif isinstance(value, nx.classes.reportviews.EdgeViewABC):
                        style[key] = [v for *e, v in value]
            elif isinstance(value, dict):
                _sanitize_leaves(value)

    def _update(style: dict, current: dict):
        for key, value in style.items():
            if key not in current:
                current[key] = value
                continue

            # Style leaves are by definition not to be recurred into
            if isinstance(value, dict) and (key not in style_leaves):
                _update(value, current[key])
            elif value is None:
                del current[key]
            else:
                current[key] = value

    old_style = copy_with_deep_values(current)

    try:
        if style is None:
            styles = []
        elif isinstance(style, (dict, str)):
            styles = [style]
        else:
            styles = list(style)

        if kwargs:
            styles.append(kwargs)

        for style in styles:
            if style == "default":
                reset()
            else:
                if isinstance(style, str):
                    current = get_style(style)
                else:
                    _sanitize_leaves(style)
                    unflatten_style(style)
                    _update(style, current)
    except:
        current = old_style
        raise


def reset() -> None:
    """Reset to default style."""
    global current
    current = copy_with_deep_values(styles["default"])


@contextmanager
def context(style: Optional[str | dict | Sequence] = None, **kwargs):
    """Create a style context for iplotx.

    Parameters:
        style: A single style specification or a list of style specifications, which are then
            applied in order. Each style can be a string (for an existing style) or a dictionary
            with the elements that are to change.
        **kwargs: Additional style changes to be applied at the end of all styles.

    Yields:
        A context manager that applies the style and reverts it back to the previous one upon exit.
    """
    current = get_style()
    try:
        use(style, **kwargs)
        yield
    finally:
        use(["default", current])


def unflatten_style(
    style_flat: dict[str, str | dict | int | float],
) -> None:
    """Convert a flat or semi-flat style into a fully structured dict.

    Parameters:
        style_flat: A flat dictionary where keys may contain underscores, which are taken to signify
            subdictionaries.

    NOTE: The dict is changed *in place*.

    Example:
        >>> style = {'vertex_size': 20}
        >>> unflatten_style(style)
        >>> print(style)
        {'vertex': {'size': 20}}
    """

    def _inner(style_flat: dict):
        keys = list(style_flat.keys())

        for key in keys:
            if "_" not in key:
                continue

            keyhead, keytail = key.split("_", 1)
            value = style_flat.pop(key)
            if keyhead not in style_flat:
                style_flat[keyhead] = {
                    keytail: value,
                }
            else:
                style_flat[keyhead][keytail] = value

        for key, value in style_flat.items():
            if isinstance(value, dict) and (key not in style_leaves):
                _inner(value)

    # top-level adjustments
    if "zorder" in style_flat:
        style_flat["network_zorder"] = style_flat["grouping_zorder"] = style_flat.pop(
            "zorder"
        )

    # Begin recursion
    _inner(style_flat)


def rotate_style(
    style: dict[str, Any],
    index: Optional[int] = None,
    key: Optional[Hashable] = None,
    props: Optional[Sequence[str]] = None,
) -> dict[str, Any]:
    """Rotate leaves of a style for a certain index or key.

    Parameters:
        style: The style to rotate.
        index: The integer to rotate the style leaves into.
        key: For dict-like leaves (e.g. vertex properties specified as a dict-like object over the
            vertices themselves), the key to use for rotation (e.g. the vertex itself).
        props: The properties to rotate, usually all leaf properties.

    Returns:
        A style with rotated leaves, which describes the properties of a single element (e.g.
        vertex).

    Example:
        >>> style = {'vertex': {'size': [10, 20]}}
        >>> rotate_style(style, index=0)
        {'vertex': {'size': 10}}
    """
    if (index is None) and (key is None):
        raise ValueError(
            "At least one of 'index' or 'key' must be provided to rotate_style."
        )

    if props is None:
        props = tuple(prop for prop in style_leaves if prop not in nonrotating_leaves)

    style = copy_with_deep_values(style)

    for prop in props:
        val = style.get(prop, None)
        if val is None:
            continue
        # Try integer indexing for ordered types
        if (index is not None) and isinstance(
            val, (tuple, list, np.ndarray, pd.Index, pd.Series)
        ):
            style[prop] = np.asarray(val)[index % len(val)]
        # Try key indexing for unordered, dict-like types
        if (
            (key is not None)
            and (not isinstance(val, (str, tuple, list, np.ndarray)))
            and hasattr(val, "__getitem__")
        ):
            # If only a subset of keys is provided, default the other ones
            # to the empty type constructor (e.g. 0 for ints, 0.0 for floats,
            # empty strings).
            if key in val:
                style[prop] = val[key]
            else:
                valtype = type(next(iter(val.values())))
                style[prop] = valtype()

    return style
