import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from collections import Iterable
from matplotlib.cbook import flatten
from itertools import cycle


def facet_plot(dframe, facets, props, ydata, layout=None, newfig=True, figsize=None,
               legend=True, individual_legends=False, hide_additional_axes=True, zorder='default', **kwargs):
    if newfig:
        nr_facets = len(dframe.groupby(facets))
        if layout is None:
            for i in range(2, nr_facets // 2):
                if nr_facets % i == 0:
                    layout = (nr_facets // i, i)
                    break
            if layout is None:
                n = int(np.ceil(nr_facets / 2))
                layout = (n, 2)
        fig, axs = plt.subplots(
            nrows=layout[0],
            ncols=layout[1],
            sharex=True, sharey=True, figsize=figsize
        )
        if hide_additional_axes:
            for ax in fig.axes[nr_facets:]:
                ax.set_axis_off()
    else:
        fig = plt.gcf()
        axs = fig.axes

    cycl = cycle(plt.rcParams['axes.prop_cycle'])
    prop_styles = {ps: next(cycl) for ps, _ in dframe.groupby(props)}

    if zorder is 'default':
        dz = 1
        zorder = 0
    elif zorder is 'reverse':
        dz = -1
        zorder = 0
    else:
        dz = 0

    if legend:
        ax0 = fig.add_subplot(111, frame_on=False, zorder=-9999)
        ax0.set_axis_off()
        plot_kwargs = kwargs.copy()
        for k in ['logx', 'logy', 'loglog']:
            plot_kwargs.pop(k, None)
        for l, p in prop_styles.items():
            ax0.plot([], label=str(l), **p, **plot_kwargs)
        ax0.legend(loc='center left', bbox_to_anchor=(1, 0.5), fontsize='x-small')
    for ax, (ps, df) in zip(flatten(axs), dframe.groupby(facets, squeeze=False)):
        for prop, df_prop in df.groupby(props):
            df_prop[ydata].plot(ax=ax, label=str(prop), zorder=zorder, **prop_styles[prop], **kwargs)
            zorder += dz
        #  ax.title(0.5, 0.1, '{},{}'.format(*ps), transform=ax.transAxes, fontsize='small')

        ax.set_title('; '.join([str(x) for x in ps]) if isinstance(ps, tuple) else str(ps), fontsize='x-small')
        if individual_legends:
            ax.legend(fontsize='x-small')

    plt.sca(ax)
    rect = (0, 0, 0.85, 1) if legend else (0, 0, 1, 1)
    plt.tight_layout(rect=rect, pad=0.1)
    return fig, axs


class CurvedText(mpl.text.Text):
    """A text object that follows an arbitrary curve."""

    def __init__(self, x, y, text, axes, **kwargs):
        super(CurvedText, self).__init__(x[0],y[0],' ', axes, **kwargs)

        axes.add_artist(self)

        # # saving the curve:
        self.__x = x
        self.__y = y
        self.__zorder = self.get_zorder()

        # # creating the text objects
        self.__Characters = []
        for c in text:
            t = mpl.text.Text(0, 0, c, **kwargs)
            # resetting unnecessary arguments
            t.set_ha('center')
            t.set_rotation(0)
            t.set_zorder(self.__zorder +1)

            self.__Characters.append((c,t))
            axes.add_artist(t)

    # # overloading some member functions, to assure correct functionality
    # # on update
    def set_zorder(self, zorder):
        super(CurvedText, self).set_zorder(zorder)
        self.__zorder = self.get_zorder()
        for c,t in self.__Characters:
            t.set_zorder(self.__zorder+1)

    def draw(self, renderer, *args, **kwargs):
        """
        Overload of the Text.draw() function. Do not do
        do any drawing, but update the positions and rotation
        angles of self.__Characters.
        """
        self.update_positions(renderer)

    def update_positions(self,renderer):
        """
        Update positions and rotations of the individual text elements.
        """

        # preparations

        # # determining the aspect ratio:
        # # from https://stackoverflow.com/a/42014041/2454357

        # # data limits
        xlim = self.axes.get_xlim()
        ylim = self.axes.get_ylim()
        # #  Axis size on figure
        figW, figH = self.axes.get_figure().get_size_inches()
        # #  Ratio of display units
        _, _, w, h = self.axes.get_position().bounds
        # # final aspect ratio
        aspect = ((figW * w)/(figH * h))*(ylim[1]-ylim[0])/(xlim[1]-xlim[0])

        # points of the curve in figure coordinates:
        x_fig,y_fig = (
            np.array(l) for l in zip(*self.axes.transData.transform([
            (i,j) for i,j in zip(self.__x,self.__y)
            ]))
        )

        # point distances in figure coordinates
        x_fig_dist = (x_fig[1:]-x_fig[:-1])
        y_fig_dist = (y_fig[1:]-y_fig[:-1])
        r_fig_dist = np.sqrt(x_fig_dist**2+y_fig_dist**2)

        # arc length in figure coordinates
        l_fig = np.insert(np.cumsum(r_fig_dist),0,0)

        # angles in figure coordinates
        rads = np.arctan2((y_fig[1:] - y_fig[:-1]),(x_fig[1:] - x_fig[:-1]))
        degs = np.rad2deg(rads)


        rel_pos = 10
        for c,t in self.__Characters:
            # finding the width of c:
            t.set_rotation(0)
            t.set_va('center')
            bbox1  = t.get_window_extent(renderer=renderer)
            w = bbox1.width
            h = bbox1.height

            # ignore all letters that don't fit:
            if rel_pos+w/2 > l_fig[-1]:
                t.set_alpha(0.0)
                rel_pos += w
                continue

            elif c != ' ':
                t.set_alpha(1.0)

            # finding the two data points between which the horizontal
            # center point of the character will be situated
            # left and right indices:
            il = np.where(rel_pos+w/2 >= l_fig)[0][-1]
            ir = np.where(rel_pos+w/2 <= l_fig)[0][0]

            # if we exactly hit a data point:
            if ir == il:
                ir += 1

            # how much of the letter width was needed to find il:
            used = l_fig[il]-rel_pos
            rel_pos = l_fig[il]

            # relative distance between il and ir where the center
            # of the character will be
            fraction = (w/2-used)/r_fig_dist[il]

            # # setting the character position in data coordinates:
            # # interpolate between the two points:
            x = self.__x[il]+fraction*(self.__x[ir]-self.__x[il])
            y = self.__y[il]+fraction*(self.__y[ir]-self.__y[il])

            # getting the offset when setting correct vertical alignment
            # in data coordinates
            t.set_va(self.get_va())
            bbox2 = t.get_window_extent(renderer=renderer)

            bbox1d = self.axes.transData.inverted().transform(bbox1)
            bbox2d = self.axes.transData.inverted().transform(bbox2)
            dr = np.array(bbox2d[0]-bbox1d[0])

            # the rotation/stretch matrix
            rad = rads[il]
            rot_mat = np.array([
                [np.cos(rad), np.sin(rad)*aspect],
                [-np.sin(rad)/aspect, np.cos(rad)]
            ])

            # # computing the offset vector of the rotated character
            drp = np.dot(dr,rot_mat)

            # setting final position and rotation:
            t.set_position(np.array([x,y])+drp)
            t.set_rotation(degs[il])

            t.set_va('center')
            t.set_ha('center')

            # updating rel_pos to right edge of character
            rel_pos += w-used