# =============================================================================
# Tools for Axes manage
# =============================================================================
import math
import numpy as np
from operator import contains
from string import Template

import matplotlib as _mpl
import matplotlib.patheffects as path_effects
import matplotlib.pyplot as _plt
import numpy as _np
import pandas as _pd
import seaborn as _sns

from ._style import StyleSW as style


class _axTools(object):
    artistList = (
        _mpl.patches.Wedge,
        _mpl.patches.Ellipse,
        _mpl.patches.Circle,
        _mpl.patches.Shadow,
        _mpl.lines.Line2D,
        _mpl.collections.PathCollection,
        _mpl.collections.PolyCollection,
    )

    def gca(ax):
        """
         Set the current axes. This is a convenience function to make it easier to use : func : ` pylab. gca ` without having to re - create the axes in the meantime.

         @param ax - the axis to set the current axes to. If you don't want to change the current axes use
        """
        if isinstance(ax, _mpl.axes._subplots.SubplotBase):
            return ax
        else:
            return _plt.gca()

    def orient(ax=None):
        """
         Orient the axes. This is a no - op if ax is None

         @param ax - axis to orient. If None a new axis is created
        """
        orient = "v"
        factor = 0
        axes = _axTools.get_axes(ax)
        rects = axes[0].patches
        if len(rects) > 0:
            for k, t in enumerate(axes):
                rects = t.patches
                if len(rects) > 1:
                    if not isinstance(rects[0], _mpl.patches.Wedge):
                        q = [round(x.get_width(), 2) for x in rects]
                        if all(x == q[0] for x in q):
                            orient = "v"
                            factor = q[0]
                        else:
                            orient = "h"
                            factor = rects[0].get_height()
                elif len(rects) == 1:
                    if not isinstance(rects[0], _mpl.patches.Wedge):
                        if rects[0].get_x() == 0:
                            orient = "h"
                        else:
                            orient = "v"
        return orient, factor

    def get_axes(ax=None):
        """
         Get or set the Axes object. This is a low - level function to allow you to get an Axes object from a list of axes or a single axes

         @param ax - If None ( default ) the current axes is returned
        """
        axes = []
        if isinstance(ax, _sns.axisgrid.FacetGrid):
            axes = ax.axes[0]
        elif isinstance(ax, list):
            axes = ax
        else:
            axes = [ax]
        return axes

    def getColorList(size: int, color):
        """
         Returns a list of colors that can be used to draw a color. The list is sorted by color and size

         @param size - the size of the color list
         @param color - the color to draw ( one of L { twisted. internet. colors. BGR }
        """
        colors = list()
        if type(color) is str:
            colors = [color] * size
        elif type(color) is list:
            miss = size - len(color)
            if miss < 0:
                colors = color[:size]
            elif miss > 0:
                colors = color
                for i in range(0, miss):
                    colors.append("white")
            else:
                colors = color
        return colors

    def values_bars(ax=None, normBy: str() = "c", loc: str() = "top", stack=False):
        """
         Plot values of a function. This is a wrapper around plot_bars that does not require matplotlib

         @param ax - axis to plot on.
         @param normBy - normalization method to use.
         @param loc - location to plot ( top bottom ). Default top
         @param stack - if True plot stack instead of normalizing. Default
        """
        pos = 0
        orient, factor = _axTools.orient(ax)
        data = _pd.DataFrame(
            columns=[
                "x",
                "y",
                "value",
                "value_norm",
                "index_color",
                "ax",
                "type",
                "color",
            ]
        )
        data = data.astype(
            {
                "x": "float32",
                "y": "float32",
                "value": "float32",
                "ax": "int32",
                "index_color": "int32",
            }
        )
        axes = _axTools.get_axes(ax)
        # Getting values
        for k, t in enumerate(axes):
            h = 0
            hist = {}
            rects = t.patches
            if len(rects) > 1:
                for rect in rects:
                    if not isinstance(rect, _mpl.patches.Wedge):
                        color = rect.get_fc()
                        if orient == "v":
                            value = (
                                0
                                if math.isnan(rect.get_height()) == True
                                else rect.get_height()
                            )
                            x = rect.get_x() + rect.get_width() / 2
                            if stack:
                                h = hist[x] if x in hist else 0
                            else:
                                h = 0
                            hist[x] = value + h
                            y = value + h
                        elif orient == "h":
                            value = (
                                0 if math.isnan(rect.get_width()) else rect.get_width()
                            )
                            y = rect.get_y() + rect.get_height() / 2
                            if stack:
                                h = hist[y] if y in hist else 0
                            else:
                                h = 0
                            hist[y] = value + h
                            x = value + h
                        data.loc[pos] = (x, y, value, 0, 0, int(k), "bar", color)
                        pos = pos + 1
            elif len(rects) == 1:
                rect = rects[0]
                if not isinstance(rect, _mpl.patches.Wedge):
                    color = rect.get_fc()
                    if orient == "h":
                        value = rect.get_width()
                        x = rect.get_width()
                        y = rect.get_x()
                    elif orient == "v":
                        value = rect.get_height()
                        y = rect.get_height()
                        x = rect.get_y()
                    data.loc[pos] = (x, y, value, 0, 0, int(k), "bar", color)

        # Grouping data by condition
        if normBy == "g":
            if isinstance(ax, _sns.axisgrid.FacetGrid):
                data.loc[:, "index_color"] = data.loc[:, "ax"]
            else:
                for i, color in enumerate(data["color"].unique()):
                    dt = data[data["color"] == color]
                    data.loc[dt.index, "index_color"] = i
        elif normBy == "c":
            if orient == "v":
                data.loc[:, "index_color"] = data.loc[:, "x"].round(0).astype("int")
            else:
                data.loc[:, "index_color"] = data.loc[:, "y"].round(0).astype("int")
        data.drop("color", axis=1, inplace=True)

        # Set position values
        vals = data.loc[:, "value"]
        vals = vals.apply(lambda x: x if x > 0 else -x)
        if loc == "center":
            vals = vals.apply(lambda x: x / 2)
        if orient == "v":
            if loc in ["bottom", "center"]:
                data.loc[:, "y"] = data.loc[:, "y"] - vals
        else:
            if loc in ["bottom", "center"]:
                data.loc[:, "x"] = data.loc[:, "x"] - vals
        return data

    def values_line(ax, orient):
        """
         Draws a line between values. This is a wrapper around : func : ` matplotlib. Axes. values ` that does not need to be implemented by subclasses

         @param ax - axis on which to draw
         @param orient - orientation of the line ( horizontal or vertical )
        """
        lnum = 0
        data = _pd.DataFrame(
            columns=["x", "y", "value", "value_norm", "index_color", "ax", "type"]
        )
        axes = _axTools.get_axes(ax)
        for l in axes[0].lines:
            if orient == "v":
                y = l.get_ydata()
                if not _np.isnan(y).any():
                    lnum += 1
                else:
                    x = l.get_xdata()
                    if not _np.isnan(x).any():
                        lnum += 1
        if lnum > 0:
            x_all = list()
            y_all = list()
            ax_all = list()
            color_all = list()

            for k, ax in enumerate(axes):
                lines = ax.lines
                if len(lines) > 0:
                    for j, l in enumerate(lines):
                        y = l.get_ydata()
                        ax_all.extend([k for x in y])
                        color_all.extend([j for x in y])
                        if all(isinstance(p, str) for p in y):
                            y = list(_np.arange(0, len(y)))
                        for i, txt in enumerate(y):
                            x = l.get_xdata()
                            if any(isinstance(p, str) for p in x):
                                x = list(_np.arange(0, len(x)))
                        if len(x) > 0 and len(y) > 0:
                            x_all.extend(x)
                            y_all.extend(y)
            # Set value location
            if orient == "h":
                data["x"] = x_all
                data["y"] = y_all
                data["value"] = x_all
            else:
                data["y"] = y_all
                data["x"] = x_all
                data["value"] = y_all
            # Set other columns values
            data["index_color"] = color_all
            data["ax"] = ax_all
            data["type"] = "line"
            data["value_norm"] = 0
        # Setting data types
        data = data.astype({"index_color": "int32", "ax": "int32"})
        data = data[_pd.notna(data["value"])].reset_index()
        return data.drop("index", axis=1)

    def values_points(ax):
        """
         Return a list of points that can be plotted. This is useful for plotting data that is not a plot but a series of plot points.

         @param ax - Axes on which to plot the data. If None the plot will be shown on the current axes
        """
        data = _pd.DataFrame(
            columns=["x", "y", "value", "value_norm", "index_color", "ax", "type"]
        )
        axes = _axTools.get_axes(ax)
        # Getting value points
        for i, h in enumerate(axes):
            d = h.collections
            for d1 in d:
                data1 = _pd.DataFrame()
                values = d1.get_offsets().tolist()
                if len(values) > 0:
                    rg = range(len(values))
                    data1["x"] = [v[0] for v in values]
                    data1["y"] = [v[1] for v in values]
                    data1["value"] = [v[1] for v in values]
                    data1["index_color"] = [0 for x in rg]
                    data1["ax"] = [i for x in rg]
                    data1["type"] = ["point" for x in rg]
                    data1["value_norm"] = 0
                    data = _pd.concat([data, data1])
        return data

    def plot_values(
        ax=None,
        data: _pd.DataFrame = None,
        minvalue: float = -np.Infinity,
        maxvalue: float = np.Infinity,
        dec: int = 2,
        fontsize=_mpl.rcParams["font.size"],
        fontweight: str = "normal",
        color="white",
        display: str = "v",
        prefix: str = "",
        postfix: str = "",
        orient: str = "v",
        **kwargs,
    ):
        """
        Plot values as a 2D plot. This is a wrapper around matplotlib. pyplot. plot_values that allows you to specify the minimum and maximum values to plot in one call.

        @param ax - matplotlib axis ( default : None ) If provided the plot will be added to this axis
        @param data - pandas DataFrame ( default : None ) If provided the plot
        @param minvalue
        @param dec
        @param fontsize
        @param fontweight
        @param color
        @param display
        @param prefix
        @param postfix
        @param orient
        """
        # Initial variables
        template = _axTools.template_print_value(display)
        va = kwargs["va"]
        ha = kwargs["ha"]
        if fontsize.isnumeric():
            fontsize = fontsize * 0.6 if fontsize > 10 else fontsize - 2
        axes = _axTools.get_axes(ax)

        shadow = kwargs["shadow"]
        shadowcolor = kwargs["shadowcolor"]
        _, kwargs = _axTools.get_init_kwargs(kwargs, get_d=False)

        # Check if is Digit function

        def isDigit(x):
            try:
                x = float(x)
                return True
            except ValueError:
                return False

        # Segment data with True values
        data = (
            data[data["value"].notnull()].copy().reset_index(inplace=False, drop=True)
        )
        colors = _axTools.getColorList(data.shape[0], color)
        # Setting data types
        if dec > 0:
            vector = data["value"].to_numpy()
            for i, row in enumerate(vector):
                if isinstance(row, (float)):
                    vector[i] = round(row, dec)
            data["value"] = vector
        else:
            data["value"] = data["value"].astype("int32")
        # Print values
        for k, g in enumerate(axes):
            data_sub = data[(data["ax"] == k)]
            for c, row in zip(colors, data_sub.to_numpy()):
                x = row[0]
                y = row[1]
                value = row[2]
                value_norm = round(row[3] * 100, dec)
                graph = row[6]

                valid_value = True
                if isDigit(value) and (minvalue is not None):
                    valid_value = (
                        True if (value >= minvalue and value <= maxvalue) else False
                    )

                if valid_value:
                    # Choosing graph type
                    if graph == "bar":
                        if orient == "v":
                            if va == "auto":
                                va_h = "bottom" if y >= 0 else "top"
                            else:
                                va_h = va
                            ha_h = "center" if ha == "auto" else ha

                        else:
                            if ha == "auto":
                                ha_h = "left" if x > 0 else "right"
                            else:
                                ha_h = ha
                            va_h = "center_baseline" if va == "auto" else va

                        text = g.text(
                            x,
                            y,
                            prefix
                            + template.substitute(value=value, value_norm=value_norm)
                            + postfix,
                            fontsize=fontsize,
                            color=c,
                            ha=ha_h,
                            va=va_h,
                            fontweight=fontweight,
                            **kwargs,
                        )
                    else:
                        ha_h = "center" if ha == "auto" else ha
                        va_h = "center_baseline" if va == "auto" else va

                        if graph == "line":
                            text = g.text(
                                x,
                                y,
                                prefix
                                + template.substitute(
                                    value=value, value_norm=value_norm
                                )
                                + postfix,
                                fontsize=fontsize,
                                color=c,
                                ha=ha_h,
                                va=va_h,
                                fontweight=fontweight,
                                **kwargs,
                            )
                        elif graph == "point":
                            text = g.text(
                                x,
                                y,
                                prefix + value + postfix,
                                fontsize=fontsize,
                                ha=ha_h,
                                va=va_h,
                                color=c,
                                **kwargs,
                            )
                    if shadow > 0:
                        text.set_path_effects(
                            [
                                path_effects.Stroke(
                                    linewidth=shadow, foreground=shadowcolor, alpha=0.8
                                ),
                                path_effects.Normal(),
                            ]
                        )

    def normalize_data(df: _pd.DataFrame = None):
        """
        Normalize data to make it easier to read. This is a no - op if data is already normalized

        @param df - pandas DataFrame
        """
        data = df.copy(deep=True)
        # Normalize data values
        grp = data.groupby(["type", "index_color"], as_index=False)["value"].sum()
        for row in data.itertuples():
            total = grp["value"][
                (grp["type"] == row.type) & (grp["index_color"] == row.index_color)
            ].values[0]
            data.loc[row.Index, "value_norm"] = data.loc[row.Index, "value"] / total
        return data

    def values_pie(ax, valid_wedges, frm, fontsize, dec, minvalue):
        """
        Plot pie chart values. This is a wrapper around plot_values.

        @param ax - Axes to plot on. Must have been created before.
        @param valid_wedges - List of valid edge positions.
        @param frm - Font to use for plot. If None defaults to font_size ( 1 ).
        @param fontsize - Font size in points. If None defaults to fontsize ( 1 ).
        @param dec - Decalization parameter passed to plot_values.
        @param minvalue - Minimum value to plot ( default 0 )
        """
        data = _pd.DataFrame()
        x = y = list()
        val_list = label_list = list()
        for i, p in enumerate(valid_wedges):
            val = p.theta2 - p.theta1
            ang = (p.theta2 - p.theta1) / 2.0 + p.theta1
            y.append(_np.sin(_np.deg2rad(ang)))
            x.append(_np.cos(_np.deg2rad(ang)) - 0.2)
            val_list.append((val / 360) * 100)
            label_list.append(p.get_label())
        data["x"] = x
        data["y"] = y
        data["val"] = val_list
        data["labels"] = label_list
        # Setting bbox style
        bbox_props = dict(
            boxstyle="round", ec=(0.1, 0.1, 0.1, 0.2), fc=(1.0, 1.0, 0.9, 0.75)
        )
        kw = dict(bbox=bbox_props, zorder=99, va="center")
        # Printing data
        for i in range(data.shape[0]):
            if data.val[i] > minvalue:
                value = frm % _np.round(data.val[i], dec) + "%"
                ax.annotate(
                    value,
                    xy=(data.x[i], data.y[i]),
                    fontsize=fontsize,
                    xytext=((0.75 * data.x[i]), (0.45 * data.y[i])),
                    **kw,
                )
        return data

    def template_print_value(display: str = str()):
        """
        @brief Prints the value of template_display. It is used to check if there is a value in template_display and if it is print it.
        @param display the value to print ( default : None )
        """
        if display == "f":
            template = Template("$value\n$value_norm%")
        elif display == "v":
            template = Template("$value")
        elif display == "p":
            template = Template("$value_norm%")
        else:
            raise ValueError(
                'The alternatives for display are "f=full", "v=values" or "p=percent"'
            )
        return template

    def get_init_kwargs(
        kwargs,
        remove_kwargs=True,
        get_d=True,
        inplace_kwargs=False,
        defaults: dict = {},
    ) -> tuple:
        """
        @brief Get kwargs to be used for initialization. This is a wrapper around : func : ` dict. get ` to allow the user to specify a set of keyword arguments that will be used to initialize the model.
        @param kwargs A dictionary of keyword arguments to be initialized.
        @param remove_kwargs If True the kwargs will be removed from the initial kwargs.
        @param get_d If True the dictionary will be returned as - is.
        @param inplace_kwargs If True the kwargs will be copied to the new dictionary instead of returned.
        @param defaults A dictionary of default values to be used for initialization.
        @return A tuple of the same length as the ` ` kwargs ` ` dictionary
        """
        keys = {
            "color": style.get_ticksColor(),
            "colors": style.get_ticksColor(),
            "grid_color": style.get_ticksColor(),
            "grid_alpha": 0.35,
            "fontweight": "normal",
            "fontsize": _mpl.rcParams["axes.labelsize"],
            "shadow": 0,
            "shadowcolor": style.get_axesColor(),
            "xpad": 0,
            "ypad": 0,
            "va": "top",
            "ha": "center",
        }

        for k, v in defaults.items():
            keys[k] = v

        dic_args = {}
        if get_d:
            values = list()
            for k, v in keys.items():
                values.append(kwargs.get(k, v))
            dic_args = dict(zip(keys, values))

        if remove_kwargs:
            for key in keys:
                if key in kwargs:
                    del kwargs[key]

        if inplace_kwargs:
            for k, v in dic_args.items():
                kwargs[k] = v

        return dic_args, kwargs
