__version__ = "0.4"
__author__ = "Thomas Baldauf"
__email__ = "thomas.baldauf@dlr.de"
__license__ = "MIT"
__birthdate__ = '15.11.2021'
__status__ = 'prod' # options are: dev, test, prod

"""
Small plotting library for bar plots and line plots. Usage is optional / voluntary...
"""

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from pylab import rcParams
import matplotlib.patches as mpatches
import matplotlib as mpl
from functools import partial


def add_value_labels(ax, spacing=.9, fmt="{0:f}", stacked=False, legend="best"):
    """Add labels to the end of each bar in a bar chart.

    Arguments:
        ax (matplotlib.axes.Axes): The matplotlib object containing the axes
            of the plot to annotate.
        spacing (int): The distance between the labels and the bars.
    """

    my_figure = ax.get_figure()
    L = abs(ax.get_ylim()[1] - ax.get_ylim()[0])
    min_height = L*0.015
    small_height = L*0.05
    # print("min_heigth", min_height)

    # For each bar: Place a label
    for rect in ax.patches:
        # Get X and Y placement of label from rect.
        outside = False

        if abs(rect.get_height()) <= min_height:
            continue

        elif abs(rect.get_height()) <= small_height:
            outside = True
            y_value = rect.get_y() + rect.get_height() * 0.1
            x_value = rect.get_x() + rect.get_width() * 1.05

        # print("rect height", rect.get_height())
        else:
            y_value = rect.get_y() + rect.get_height() * 0.41
            x_value = rect.get_x() + rect.get_width() * 0.50

        # print("yvalue",x_value,y_value)
        # Number of points between bar and label. Change to your liking.

        space = spacing
        # Vertical alignment for positive values
        va = 'bottom'

        # If value of bar is negative: Place label below bar
        if y_value < 0:
            # Invert space to place label below
            space *= -1
            # Vertically align label at top
            va = 'top'

        # Use Y value as label and format number with one decimal place
        # label = "{:.1f}".format(y_value)
        hgt = rect.get_height()
        if hgt != 0.0:
            label = fmt.format(hgt)
        else:
            label = ""

        # Create annotation
        if outside:
            rot = 33
            ha = "left"
        else:
            rot = 0
            ha = "center"

        t = ax.annotate(
            label,  # Use `label` as label
            (x_value, y_value),  # Place label at end of the bar
            xytext=(0, space),  # Vertically shift label by `space`
            textcoords="offset points",  # Interpret `xytext` as offset in points
            ha=ha,  # Horizontally center label
            va=va)  # Vertically align label differently for

        if not outside:
            t.set_bbox(dict(facecolor='white', alpha=0.175, edgecolor='white',linewidth=0.0))
        # positive and negative values.


def matplotlib_barplot(data, xlabel, ylabel, title, color="indigo", size=(6, 5), tight_layout=True,
                       fmt="{0:f}", stacked=False, show_labels=True, legend="best", show=True,yerr=None,ax=None):
    """
    Creates a bar plot in matplotlib. The 'macro' plotting style will be used if it is set up

    :param xlabel: x axis  label
    :param ylabel: y axis  label
    :param title: plot title
    :param color: plot color
    :param size: figure size
    :param tight_layout: re-format plot
    :param fmt: number format, {0:d} or {0:f} for example
    :param stacked: stack column? or show side-by side (boolean)
    :paarm show_labels: show the column labels at bottom? (boolean)
    :param legend: location of legend, or 'off'
    :param show: if False, only the figure is returned. Else plot is shown
    :param ax: axis to plot on
    :param yerr: data for error bar (if any)
    """

    # if legend is 'off', no legend is shown

    plt.rcParams['axes.ymargin'] = .4
    rcParams['figure.figsize'] = size

    try:
        plt.style.use("macro") # if 'macro' plotting style is installed -> use it!
    except:
        pass

    plt.grid(alpha=0.0)

    if not isinstance(data, pd.DataFrame):
        if ylabel is None:
            ylabel = "Data"
        data = pd.DataFrame({ylabel: data})
        data.set_index(ylabel)

    # plt.figure(figsize=size)
    if ax is None:
        ax = plt.gca()
    data.plot.bar(color=color, stacked=stacked,ax=ax,yerr=yerr,capsize=6,ecolor="black")
    plt.gca().set_title(title)
    # plt.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=True)
    plt.ylabel(ylabel)
    plt.xlabel(xlabel)

    if show_labels:
        add_value_labels(plt.gca(), fmt=fmt)

    if legend != "off":
        plt.legend(loc=legend)
    else:
        old_legend = plt.gca().get_legend()
        if old_legend is not None:
            old_legend.remove()
    if tight_layout:
        plt.tight_layout()

    # fix reverse labeling
    # https://stackoverflow.com/questions/46908085/bar-plot-does-not-respect-order-of-the-legend-text-in-matplotlib
    #
    ax = plt.gca()
    h, l = ax.get_legend_handles_labels()
    if legend != "off":
        ax.legend(h[::-1], l[::-1],loc=legend,framealpha=0.0)

    if show:
        plt.show()
    else:
        return plt.gcf()


def matplotlib_lineplot(data, xlabel=None, ylabel=None, title="", xlim=None, ylim=None, color="indigo", legend="best",marker=None,show=False,ax=None):
    """
    creates a line plot in matplotlib. The 'macro' plotting style will be used if it is set up

    :param xlabel: x axis  label
    :param ylabel: y axis  label
    :param title: plot title
    :param color: plot color
    :param xlim, ylim: tuples of plot limits
    :param legend: location of legend, or 'off', or 'outside'
    :param show: if False, only the figure is returned. Else plot is shown
    :param ax: axis to plot on

    when overlapping barplot and lineplot, make sure you add an additional axis, see
    https://stackoverflow.com/questions/42948576/pandas-plot-does-not-overlay

    """
    try:
        plt.style.use("macro")
    except:
        pass


    if not isinstance(data, pd.DataFrame):
        if ylabel is None:
            label = "Data"
        data = pd.DataFrame({ylabel: data})
        data.set_index(ylabel)

    if ax is None:
        ax = plt.gca()

    for col in data.columns:
        ax.plot(data[col],color=color,marker=marker,label=col)

    plt.xlim(xlim)
    plt.ylim(ylim)
    plt.title(title)

    if legend != "off":
        plt.legend(loc=legend)
    elif legend == "outside":
        plt.legend(loc=(1.02,0))

    elif legen == "off":
        old_legend = plt.gca().get_legend()
        if old_legend is not None:
            old_legend.remove()

    if ylabel is not None:
        plt.ylabel(ylabel)

    if xlabel is not None:
        plt.gca().set_xlabel(xlabel)

    if ylabel is not None:
        plt.gca().set_ylabel(ylabel)
    plt.rcParams['axes.ymargin'] = .4

    if show:
        plt.show()
    else:
        return plt.gcf()





class Point:
    """
    a rudimentary point in 2d space
    """

    def __init__(self, x, y):
        self._x = x
        self._y = y
        self.offset_y = 0 # offset for bands

    @property
    def x(self):
        return self._x

    @property
    def y(self):
        return self._y + self.offset_y

class Label:
    """
    a rudimentary node label
    """

    names = []

    def __init__(self,x,y,text):
        self.x = x
        self.y = y
        self.text = text
        if text in self.__class__.names:
            raise RuntimeError("Name already given")
        else:
            self.__class__.names.append(self.text)

    def shift_up(self):
        # shift the label upwards a bit
        self.y += 1.8
        self.x -= 0.2 # re-center


def casteljau(t,b0,b1,b2,b3):
    """
    casteljau bezier spline for 4 points.

    :param b0-b3: the point of the bezier curve as [x,y] array
    :param t: parameter for position on the curve, float
    """
    b = np.multiply((-b0 + 3*b1 - 3*b2 + b3),np.power(t,3))  + np.multiply((3*b0 - 6*b1 + 3*b2),np.power(t,2)) + np.multiply((-3*b0 + 3*b1),t) + b0
    return b


def bezier(pstart : Point, pend: Point ,npoints = 100, curv=None):
    """
    compute the path of a 1d bezier curve

    :param pstart: starting point of the bezier curve
    :param pend: end point of the bezier curve
    :param npoints: number of discrete points to retrun
    :param curv: optionally you can give a curvature parameter here (float) default None

    :return: x,y coordinate tuple as ndarray of shape (npoints,)
    """

    # de-casteljau b-spline
    tarr = np.linspace(0,1,npoints) # parameter of the curve position
    # define 4 points: b0, b1, b2, b3

    if pstart.x < pend.x: # pstart left of pend?
        pleft = pstart
        pright = pend
    else:
        pleft = pend
        pright = pstart

    if curv is None:
        curv = .27*(pright.x-pleft.x)

    b0 = np.array([pleft.x,pleft.y])
    b1 = np.array([pleft.x+curv,pleft.y])
    b2 = np.array([pright.x-curv,pright.y])
    b3 = np.array([pright.x,pright.y])

    spline = partial(casteljau,b0=b0,b1=b1,b2=b2,b3=b3)

    b = np.array([spline(t) for t in tarr])

    return b[:,0],b[:,1]


def plot_band(pstart: Point, pend: Point, npoints=100,curv=None, weight=2):
    """
    computes a band of two bezier curves

    :param pstart: starting point of the bezier curve
    :param pend: end point of the bezier curve
    :param npoints: number of discrete points to retrun
    :param curv: optionally you can give a curvature parameter here (float) default None
    :param weight: width of the band

    :return: x,y1,y2 coordinate 3-tuple of the band
    """

    w = weight # width of the band
    d = 0 # horizontal shift of the bands realtive to each other
    e = 0 # distance of  '<' and '>' arc

    qstart1 = Point(pstart.x-d,pstart.y)
    qstart2 = Point(pstart.x-d,pstart.y)

    qend  = Point(pend.x-d,pend.y)
    qend1  = Point(pend.x-d,pend.y)
    qend2  = Point(pend.x-d,pend.y)

    pstart1 = Point(pstart.x-d+.2*e,pstart.y)
    pend1 = Point(pend.x-d-e,pend.y)

    pstart2 = Point(pstart.x+d+.2*e,pstart.y+w)
    pend2 = Point(pend.x+d-e,pend.y+w)

    bezier_start1 = bezier(qstart1,pstart1)
    bezier_start2 = bezier(qstart2,pstart2)

    bezier_end1 = bezier(pend1,pend1)
    bezier_end2 = bezier(pend2,pend2)

    bezier1 = bezier(pstart1,pend1)
    bezier2 = bezier(pstart2,pend2)

    #x = np.array(list(bezier_start1[0])+list(bezier1[0])+list(bezier_end1[0]))
    #y1 = np.array(list(bezier_start1[1])+list(bezier1[1])+list(bezier_end1[1]))
    #y2 = np.array(list(bezier_start2[1])+list(bezier2[1])+list(bezier_end2[1]))

    x = np.array(list(bezier1[0]))  # + list(bezier_end1[0]))
    y1 = np.array(list(bezier1[1])) #  + list(bezier_end1[1]))
    y2 = np.array(list(bezier2[1])) # + list(bezier_end2[1]))


    return x, y1, y2


def draw_band(fig,pstart:Point, pend:Point, color="blue",alpha=.9,weight=2,norm=1.0,show_label=True,label=None,dv=0.6,rotation=0):
    """
    draw a band graphically

    """
    x,y1,y2 = band(pstart,pend,weight=weight)
    plt.fill_between(x,y1,y2,color=color,figure=fig,alpha=alpha)
    
    my_bez_x, my_bez_y = bezier(Point(pstart.x, pstart.y), Point(pend.x, pend.y))
    
    plt.plot(my_bez_x, my_bez_y, color="black",alpha=.7) # "black")

    if show_label and label is not None:
        # draw bezier
        
        # dv = 0.6
        bx = my_bez_x[int(dv*len(my_bez_x))]
        by = my_bez_y[int(dv*len(my_bez_y))]

        
        #bx2,by2 = casteljau(t=0.55,b0=np.array([pstart.x,pstart.y]),b1=b1,b2=b2,b3=np.array([pend.x,pend.y]))
        #bx3,by3 = casteljau(t=0.66,b0=np.array([pstart.x,pstart.y]),b1=b1,b2=b2,b3=np.array([pend.x,pend.y]))
        #rot = (90-.5*np.arctan2((bx3-bx2),(by3-by2))*180.0/np.pi)%360.0 # compute slope at point
        #rot = 0

        bbox = {'pad': 2, 'facecolor': color,'edgecolor':color,'linewidth':.0}
        props = {'ha': 'center', 'va': 'center', 'bbox': bbox}

        if label is not None:
            plt.gca().text(bx,by - .0, label, props, rotation=rotation,color="white", alpha=alpha)
        else:
            plt.gca().text(bx,by , "%.2f" %(weight/norm), props, rotation=rotation,color="black",alpha=alpha)

        
        # plt.annotate(str(weight),(bx,by),color="gray",alpha=alpha)

    pend.offset_y += 1.5*weight
    pstart.offset_y += 1.5*weight

def plot_sankey(data,title="",show_plot=True,
                show_values=True,
                colors=None,
                label_rot=65,
                fontsize=10,
                separation=0.8,
                norm_factor=1.5,
                min_width=0.1,
                dy=10,
                dx=55,
                dx_space=8.5,
                alpha=0.5,
                filling_fraction = 0.85):
    """
    plots a sankey diagram from data

    :param data: list of pandas dataframes, of the following format. length at least two.
    :param title: title of the plot
    :param show_plot: boolean switch to show plot window (default True). If False, figure is returned
    :param show_values: boolean switch to print numerical values of data as text label (default True)
    :param colors: (optional) a list of colors (if None, default colors are chosen)
    :param label_rot: rotation of the labels (default 65 degrees)
    :param fontsize: size of label font
    :param separation: separation point of the label along the bands, between 0 and 1,
    :param norm_factor: normalization factor for width 
    :param min_width: minimum width of bands
    :param dy: vertical distance of bands
    :param dx: horizontal distance of bands
    :param dx_space: space between layers
    :param filling_fraction: fraction by which the bands align (scaling with dy)

    :return fig: None or matplotlib figure

    Example: 

    +----------+----------+----------+-------------+
    |  from    |   to     |  value   | color_id    |
    +----------+----------+----------+-------------+
    |    A     |  C       | 1.0      |    0        |
    +----------+----------+----------+-------------+
    |    A     |  D       | 2.0      |    1        |
    +----------+----------+----------+-------------+
    |    B     |  D       | 3.0      |    2        |
    +----------+----------+----------+-------------+

    +----------+----------+----------+-------------+
    |  from    |   to     |  value   | color_id    |
    +----------+----------+----------+-------------+
    |    C     |  E       | 5.0      |    1        |
    +----------+----------+----------+-------------+
    |    C     |  F       | 6.0      |    2        |
    +----------+----------+----------+-------------+
    |    D     |  G       | 8.0      |    3        |
    +----------+----------+----------+-------------+

    """

    Label.names = []

    if not isinstance(data,list):
        data = [data] # convert to list

    if colors is None:
        colors = [
            "lightgray",
            "skyblue",
            "wheat",
            "lightgreen",
            "lavender",
        ]
    
    # 1. normalize the values to reasonable width of bands

    # norm_factor = 1.5 #normalization factor for width

    maxval = 0.0
    maxrows = 2

    for layer in data:

        maxrows = max(maxrows,len(list(layer.iterrows())))
        for idx,row in layer.iterrows():
            val = float(row["value"])

            maxval = max(val,maxval)

    # norm_factor = 1.5
    if maxval > 0.0:
        norm_factor = 1.0 /maxval

    fig = plt.figure(figsize=(2.4*len(data),1.05*maxrows)) # 1.8*maxrows

    # 2. plot the bezier patches

    i = 10  #
    j0 = 6 # < starting position (i,j)

    # dy = 10 # 8  # vertical distance of bands
    # dx = 55 #   # horizontal distance of layers

    # dx_space = 2.5 # white space between layers

    # define some points for the unique nodes
    my_points = {} # keys will be names, value will be Points
    my_labels = {} # keys will be names, value will be Labels
    my_values = {} # labels for values

    last_j = j0

    # iterate through the dataframes
    # heights = defaultdict(lambda: 0 )

    
    class Band:
    

            instances = [] 
            def __init__(self,n_from,n_to,w,color="steelblue"):
                self.n_from = n_from 
                self.n_to = n_to 
                self.shift_y =0.0 
                self.width =  2* w
                self.osy1 = 0.0
                self.osy2 = 0.0
                self.color = color
                self.__class__.instances.append(self) 

            def plot(self):
                fromnode = self.n_from 
                tonode = self.n_to
                w = self.width
                osy1 = self.osy1
                osy2 = self.osy2
                
                #vals_lower = [fromnode.y+osy1,tonode.y+osy2]
                #vals_upper = [fromnode.y+osy1 + w,tonode.y+osy2 + w]

                point_lower1 = Point(fromnode.x,fromnode.y+osy1)
                point_lower2  = Point(tonode.x, tonode.y + osy2)
                # point_lower2 = Point(fromnode.x,fromnode.y+osy1+w)

                # make the curves more smooth 
                x, vals_lower, vals_upper = plot_band(point_lower1, point_lower2,weight=w)
                plt.fill_between(x, vals_lower, vals_upper,color=self.color)
                
                
            def maxy(self):
                w = 2 * self.width 
                return max([self.n_from.y+ self.osy1 + w, self.n_to.y+ self.osy2 + w])
            

   

    from collections import defaultdict
    bands = defaultdict(lambda: [])

    
    for indx, layer in enumerate(data):

        j =  j0 

        k = 0
        for name in layer["from"].unique():

            if name not in my_points:
                new_point = Point(i+dx_space,j)
                my_points[name] = new_point

                j += dy

            if name not in my_labels:
                my_labels[name] = Label(new_point.x-5.2,new_point.y-.2,name) # plt.annotate(name,(new_point.x-1.8, new_point.y-.1))
            else:

                my_labels[name].shift_up()

            k+=1

        j = j0

        k = 0
        for name in layer["to"].unique():

            if name not in my_points:
                new_point = Point(i+dx,j)
                my_points[name] = new_point #Point(i+dx,j)

                j += dy 

            if name not in my_labels:
                my_labels[name] = Label(new_point.x+.5*dx_space,new_point.y-.2,name) # plt.annotate(name,(new_point.x+.5*dx_space,new_point.y-.1))
            else:
                my_labels[name].shift_up()
            
            k += 1

        last_j = j - dy

        for idx, row in layer.iterrows():

            fromnode = my_points[row["from"]]
            tonode = my_points[row["to"]]
            color_idx = row["color_id"]

            my_band = Band(fromnode, tonode, w = float(row["value"])*norm_factor * dy * (1-filling_fraction), color = colors[color_idx%len(colors)])
            
            bands[fromnode].append(my_band)
            bands[tonode].append(my_band)

        i += dx

        #for name,point in my_points.items():
        #    point.offset_y = 0

        # iterate through the dataframe rows


        for point, b in bands.items():
             
            offset1 = 0.0 
            offset2 = 0.0 

            for bi in b:

                if point == bi.n_from:
                    bi.osy1 = offset1
                    offset1 += bi.width 

                elif point == bi.n_to:
                    bi.osy2 = offset2
                    offset2 += bi.width 
          

        # plot all bands 
        for band in Band.instances:
            band.plot()

        """
        for idx, row in layer.iterrows():

            color_idx = int(row["color_id"]) # index for color map

            fromnode = my_points[row["from"]]
            tonode = my_points[row["to"]]
            weight = float(row["value"])*norm_factor # 8.5e-2 #  max(min_width,float(row["value"])*norm_factor) # <- make bars thicker

            draw_band(fig,fromnode,tonode,colors[color_idx%len(colors)],weight=weight,norm=norm_factor,show_label=show_values,rotation=0,dv=separation,label=str(row["value"]))

            # update weight for next layer
        """


    for label in my_labels.values(): # draw the labels here...
        if len(label.text) <= 2:
            plt.annotate(label.text.replace("$[source]","").replace("$[sink]",""),(label.x,label.y))
        else: # rotate label for better visibility
            plt.gca().text(label.x,label.y, label.text.replace("$[source]","").replace("$[sink]",""), rotation=label_rot,fontsize=fontsize)

    # 3. return the figure or show window

    plt.axis("off")

    max_x = 0
    max_y = 0

    for point in my_points.values():
        max_x = max(max_x,point.x)
        max_y = max(max_y,point.y)

    for label in my_labels.values():
        max_x = max(max_x,label.x)
        max_y = max(max_y,label.y)

    maxy = max([b.maxy() for b in Band.instances])
    plt.xlim([3,max_x+5])
    plt.ylim([3,maxy+3])

    #plt.autoscale()

    plt.tight_layout()

    plt.title(title)

    plt.tight_layout()

    if show_plot:
        plt.show()

    return fig
