import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
from os.path import exists, basename,dirname
from icecream import ic as info
from roux.viz.ax_ import *

def plot_range(df00: pd.DataFrame,
               colvalue: str,
               colindex: str,
               k: str,
               headsize: int=15,
               headcolor: str='lightgray',
    ax: plt.Axes=None,
    ) -> plt.Axes:
    """Plot range/intervals e.g. genome coordinates as lines.

    Args:
        df00 (pd.DataFrame): input data.
        colvalue (str): column with values.
        colindex (str): column with ids.
        k (str): subset name.
        headsize (int, optional): margin at top. Defaults to 15.
        headcolor (str, optional): color of the margin. Defaults to 'lightgray'.
        ax (plt.Axes, optional): `plt.Axes` object. Defaults to None.

    Returns:
        plt.Axes: `plt.Axes` object.
    """
    df00['rank']=df00[colvalue].rank()
    x,y=df00.rd.filter_rows({colindex:k}).iloc[0,:]['rank'],df00.rd.filter_rows({colindex:k}).iloc[0,:][colvalue]
    if ax is None:
        fig,ax=plt.subplots(figsize=[1,1])
    ax=df00.set_index('rank').sort_index(0)[colvalue].plot.area(ax=ax)
    ax.annotate('', xy=(x, y),  xycoords='data',
                xytext=(x, ax.get_ylim()[1]), textcoords='data',
                arrowprops=dict(facecolor=headcolor, shrink=0,
                               width=0,ec='none',
                               headwidth=headsize,
                               headlength=headsize,
                               ),
                horizontalalignment='right', verticalalignment='top',
                )
    d_=get_axlims(ax)
    ax.text(x,y+(d_['y']['len'])*0.25,int(y),#f"{y:.1f}",
                    # transform=ax.transAxes,
                    va='bottom',ha='center',
                   )
    ax.text(0.5,0,colvalue,
                    transform=ax.transAxes,
                    va='top',ha='center',
                   )
    ax.axis(False)
    return ax
    
def plot_connections(
    dplot: pd.DataFrame,
    label2xy: dict,
    colval: str='$r_{s}$',
    line_scale: int=40,
    legend_title: str='similarity',
    label2rename: dict=None,
    element2color: dict=None,
    xoff: float=0,
    yoff: float=0,
    rectangle: dict={'width':0.2,'height':0.32},
    params_text: dict={'ha':'center','va':'center'},
    params_legend: dict={'bbox_to_anchor':(1.1, 0.5),
    'ncol':1,
    'frameon':False},
    legend_elements: list=[],
    params_line: dict={'alpha':1},
    ax: plt.Axes=None,
    test: bool=False
    ) -> plt.Axes:
    """Plot connections between points with annotations.

    Args:
        dplot (pd.DataFrame): input data.
        label2xy (dict): label to position.
        colval (str, optional): column with values. Defaults to '{s}$'.
        line_scale (int, optional): line_scale. Defaults to 40.
        legend_title (str, optional): legend_title. Defaults to 'similarity'.
        label2rename (dict, optional): label2rename. Defaults to None.
        element2color (dict, optional): element2color. Defaults to None.
        xoff (float, optional): xoff. Defaults to 0.
        yoff (float, optional): yoff. Defaults to 0.
        rectangle (_type_, optional): rectangle. Defaults to {'width':0.2,'height':0.32}.
        params_text (_type_, optional): params_text. Defaults to {'ha':'center','va':'center'}.
        params_legend (_type_, optional): params_legend. Defaults to {'bbox_to_anchor':(1.1, 0.5), 'ncol':1, 'frameon':False}.
        legend_elements (list, optional): legend_elements. Defaults to [].
        params_line (_type_, optional): params_line. Defaults to {'alpha':1}.
        ax (plt.Axes, optional): `plt.Axes` object. Defaults to None.
        test (bool, optional): test mode. Defaults to False.

    Returns:
        plt.Axes: `plt.Axes` object.
    """
    import matplotlib.patches as mpatches
    from matplotlib.collections import PatchCollection
    label2xy={k:[label2xy[k][0]+xoff,label2xy[k][1]+yoff] for k in label2xy}
    dplot['index xy']=dplot['index'].map(label2xy)
    dplot['column xy']=dplot['column'].map(label2xy)
    
    ax=plt.subplot() if ax is None else ax
    from roux.viz.ax_ import set_logos,get_subplot_dimentions
    patches=[]
    label2xys_rectangle_centers={}
    for label in label2xy:
        xy=label2xy[label]
        rect = mpatches.Rectangle(xy, **rectangle, fill=False,fc="none",lw=2,
                                  ec=element2color[label] if not element2color is None else 'k',
                                 zorder=0)

        patches.append(rect)
        line_xys=[np.transpose(np.array(rect.get_bbox()))[0],np.transpose(np.array(rect.get_bbox()))[1][::-1]]
        label2xys_rectangle_centers[label]=[np.mean(line_xys[0]),np.mean(line_xys[1])]
        inset_width=0.2
        inset_height=inset_width/get_subplot_dimentions(ax)[2]
        axin=ax.inset_axes([*[l-(off*0.5) for l,off in zip(label2xys_rectangle_centers[label],[inset_width,inset_height])],
                            inset_width,inset_height])
        if not test:
            axin=set_logos(label=label,element2color=element2color,ax=axin,test=test)
        axin.text(np.mean(axin.get_xlim()),np.mean(axin.get_ylim()),
                 label2rename[label] if not label2rename is None else label,
                  **params_text,
                 )
    dplot.apply(lambda x: ax.plot(*[[label2xys_rectangle_centers[x[k]][0] for k in ['index','column']],
                                  [label2xys_rectangle_centers[x[k]][1] for k in ['index','column']]],
                                  lw=(x[colval]-0.49)*line_scale,
                                  linestyle=params_line['linestyle'],
                                  color='k',zorder=-1,
                                  alpha=params_line['alpha'],
                                ),axis=1)            
    if params_line['annot']:
        def set_text_position(ax,x):
            xs,ys=[[label2xys_rectangle_centers[x[k]][i] for k in ['index','column']] for i in [0,1]]
            xy=[np.mean(xs),np.mean(ys)]
            if np.subtract(*xs)==0 or np.subtract(*ys)==0:
                ha,va='center','center'
                rotation=0
            else:
                if np.subtract(*xs)<0:      
                    ha,va='right','bottom'
                    xy[1]=xy[1]+0.025
                    rotation=-45
                else:
                    ha,va='right','top'
                    xy[1]=xy[1]-0.025
                    rotation=45
            ax.text(xy[0],xy[1],f"{x[colval]:.2f}",
                    ha=ha,va=va,
                    color='k',rotation=rotation,
                   bbox=dict(boxstyle="round",
                   fc='lightgray',ec=None,)
                   )
            return ax
            
        dplot.apply(lambda x: set_text_position(ax,x),axis=1)            
    from matplotlib.lines import Line2D
    legend_elements=legend_elements+[Line2D([0], [0], color='k', linestyle='solid', lw=(i-0.49)*line_scale, 
                                alpha=params_line['alpha'],
                                label=f' {colval}={i:1.1f}') for i in [1.0,0.8,0.6]]
    ax.legend(handles=legend_elements,
              title=legend_title,**params_legend)
    ax.set(**{'xlim':[0,1],'ylim':[0,1]})
    if not test:
        ax.set_axis_off()      
    return ax

def plot_kinetics(
    df1: pd.DataFrame,
    x: str,
    y: str,
    hue: str,
    cmap: str='Reds_r',
    ax: plt.Axes=None,
    test: bool=False,
    kws_legend: dict={},
    **kws_set,
    ) -> plt.Axes:
    """Plot time-dependent kinetic data.

    Args:
        df1 (pd.DataFrame): input data.
        x (str): x column.
        y (str): y column.
        hue (str): hue column.
        cmap (str, optional): colormap. Defaults to 'Reds_r'.
        ax (plt.Axes, optional): `plt.Axes` object. Defaults to None.
        test (bool, optional): test mode. Defaults to False.
        kws_legend (dict, optional): legend parameters. Defaults to {}.

    Returns:
        plt.Axes: `plt.Axes` object.
    """
    from roux.viz.ax_ import rename_legends
    from roux.viz.colors import get_ncolors
    df1=df1.sort_values(hue,ascending=False)
    info(df1[hue].unique())
    if ax is None: fig,ax=plt.subplots(figsize=[2.5,2.5])
    label2color=dict(zip(df1[hue].unique(),get_ncolors(df1[hue].nunique(),
                                                            ceil=False,
                                                            cmap=cmap,
                                                                )))
    df2=df1.groupby([hue,x],sort=False).agg({c:[np.mean,np.std] for c in [y]}).rd.flatten_columns().reset_index()
    d1=df1.groupby([hue,x],sort=False,as_index=False).size().groupby(hue)['size'].agg([min,max]).T.to_dict()
    d2={str(k):str(k)+'\n'+(f"(n={d1[k]['min']})" if d1[k]['min']==d1[k]['max'] else f"(n={d1[k]['min']}-{d1[k]['max']})") for k in d1}
    if test:info(d2)
    df2.groupby(hue,sort=False).apply(lambda df: df.sort_values(x).plot(x=x,
                                                            y=f"{y} mean",
                                                            yerr=f"{y} std",
                                                                        elinewidth=0.3,
                                                            label=df.name,
                                                            color=label2color[df.name],
                                                            lw=2,
                                                           ax=ax))
    ax=rename_legends(ax,replaces=d2,title=hue,
                     **kws_legend)
    ax.set(**kws_set)
    return ax


## plot data shape changes
def plot_steps(df1,
               coln='n',
               ax=None,
               test=False):
    """
    changes in numbers
    """
    df1=df1.sort_values(coln,ascending=False)
    df1['% change']=df1[coln].pct_change() * 100
    df1['y']=range(len(df1))
    if ax is None:
        fig,ax=plt.subplots(figsize=[4,len(df1)])
    kws_line=dict( marker='o',mfc='w',color='gray',ms=17)
    df1.iloc[:-1,:].apply(lambda x: ax.plot([0,0],[x['y'],x['y']+1],**kws_line),axis=1)
    df1.apply(lambda x: ax.text(0.005,x['y'],s=x.name,ha='left',va='center'),axis=1)
    df1.apply(lambda x: ax.text(0,x['y'],s=f"{x['y']:.0f}",ha='center',va='center'),axis=1)
    df1.apply(lambda x: ax.text(0.005,x['y']+0.33,s=f"n={x[coln]:.0f}",ha='left',va='center',
                               alpha=0.75),axis=1)
    from rohan.lib.plot.colors import saturate_color
    df1.apply(lambda x: ax.text(0.005,x['y']+0.66,s=f"{x['% change']:.1f}%" if not pd.isnull(x['% change']) else '',ha='left',va='center',
                               color=saturate_color("#FF0000",0.5+(x['% change']/-100)) if x['% change']<0 else 'g',
                               alpha=0.75),axis=1)
    ax.set(xlim=[-0.005,0.1],
          ylim=[len(df1)+0.5,-1.5])
    if not test:ax.axis('off')
    return ax