# Standard Library imports
import arrow
from datetime import datetime
from itertools import cycle
import pandas as pd
import numpy as np
from pprint import pprint as pp
import statistics
import sys
import time
import traceback
import warnings
import logging
warnings.simplefilter(action='ignore', category=pd.errors.PerformanceWarning)

# Third party imports
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import matplotlib.dates as dates

#Bokeh imports
from bokeh.io import push_notebook, show, output_notebook
from bokeh.plotting import figure, output_file, show as bokeh_show
from bokeh.models import LinearAxis, Range1d, DataRange1d, HoverTool
from bokeh.core.validation import check_integrity
from bokeh.layouts import column, row, gridplot, layout
from bokeh.models import ColumnDataSource, Div, Span, CustomJS

# Load Application imports
from dmyplant2.dReliability import demonstrated_reliability_sr
import dmyplant2

def cvset(mp, dset):
    vset = [d for col in [rec['col'] for rec in dset] for d in col]
    di = mp.get_dataitems()
    vset = [d for d in vset if not di.loc[di.name.str.fullmatch(d)].empty]
    return vset

def cplotdef(mp, lfigures):
    for f in lfigures:
        #print(f, end=', ')
        plotdef = list(lfigures.keys())
        vset = []
        for p in plotdef:
            dset = lfigures[p]
            vset += cvset(mp, dset)
        vset = list(set(vset))
    return plotdef, vset

def equal_adjust(dset, ddata, do_not_adjust=[], debug=False, qmin=0.02, qmax=0.95, minfactor=0.9, maxfactor=1.1):

    def debug_dset(dset):
        for row in dset:
            print(f"{str(row['col']):40} {row['ylim'] if 'ylim' in row else 'no ylim'}")

    data = ddata.fillna(0)
    vset = [d for col in [rec['col'] for rec in dset] for d in col]
    do_not_adjust = [x if type(x)==int else vset.index(x) for x in do_not_adjust]
    num_axes = len(dset) - len(do_not_adjust)
    dncm = [len(dset) + d for d in do_not_adjust if  d < 0]
    dncp = [d for d in do_not_adjust if  d >= 0]
    dnc = dncm + dncp
    if debug:
        print(f"nuber of data columns: {len(dset)}, number of axes:{num_axes}, Rows excluded from adjustement:{dnc}")
    k = 0
    for i, ax in enumerate(dset):
        if not i in dnc:
            for j, d in enumerate(ax['col']):
                lmin = data[d].quantile(q=qmin) if j==0 else min(lmin, data[d].quantile(q=qmin))
                lmax = data[d].quantile(q=qmax) if j==0 else max(lmax, data[d].quantile(q=qmax))
                lmin = minfactor * lmin if lmin > 0.0 else maxfactor * lmin
                lmax = minfactor * lmax if lmax < 0.0 else maxfactor * lmax
                if debug:
                    print(f"{i} {d:20} min={lmin:8.2f}, max={lmax:8.2f}")
                smax = k*(lmax - lmin) + lmax
                smin = lmin - (num_axes - (k + 1))*(lmax - lmin)
                if (smax - smin) < 1.0:
                    smin = smin - 5.0 - k; smax = smax + 5.0 - k
                ax['ylim'] = (smin, smax)
            k += 1
    if debug:
        print("\nCalculated Axes:")
        debug_dset(dset)
    return dset

def count_columns(dset):
    cnt = 0
    for c in dset:
        cnt += len(c['col'])
    return cnt

def _idx(n, s, e, x):
    return int(n * (x - s) / (e - s)+1)

def add_vlines(lines, ax, *args, **kwargs):
    for l in lines:
        ax.axvline(l, *args, **kwargs)

def add_dbokeh_vlines(lines, fig, *args, **kwargs):
    for l in lines:
        fig.add_layout(Span(location=l,dimension='height', *args, **kwargs))

def add_dbokeh_hlines(lines, fig, *args, **kwargs):
    for l in lines:
        fig.add_layout(Span(location=l,dimension='width', *args, **kwargs))

def add_table(summary, ax, *args, **kwargs):
    """
    available options for loc:
    best, upper right, upper left, lower left, lower right, center left, center right
    lower center, upper center, center, top right,top left, bottom left, bottom right
    right, left, top, bottom
    """
    ax.table(
        cellText=summary.values,
        colWidths=[0.1]*len(summary.columns),
        colLabels=summary.columns,
        cellLoc='center',
        rowLoc='center',
        *args, **kwargs)
        #loc='upper left')

def _plot(idf, x12='datetime', y1 = ['Various_Values_SpeedAct'], y2 = ['Power_PowerAct'], ylim2=(0,5000), *args, **kwargs):
    ax = idf[[x12] + y1].plot(
    x=x12,
    y=y1,
    kind='line',
    grid=True,
    *args, **kwargs)

    ax2 = idf[[x12] + y2].plot(
    x=x12,
    y=y2,
    secondary_y = True,
    ax = ax,
    kind='line',
    grid=True,
    *args, **kwargs)

    ax2.set_ylim(ylim2)
    return ax, ax2, idf

def dbokeh_chart(source, pltcfg, x='datetime', x_ax_unit=None, title=None, grid=True, legend=True, style='line', x_range=None, y_range=None, notebook=True, figsize=(8,6), *args, **kwargs):
    """wrapper function for bokeh_chart from Johannes"""
    if notebook: output_notebook(hide_banner=True)
    if title: title = str(title)
    for col in pltcfg:
        if not 'unit' in col: col['unit'] = ''
    source = ColumnDataSource(source)
    fig = bokeh_chart(source, pltcfg, x, x_ax_unit, title, grid, legend, style, x_range, y_range, figsize, *args, **kwargs)
    return fig

def bokeh_chart(source, pltcfg, x_ax='datetime', x_ax_unit=None, title=None, grid=True, legend=True, style='line', x_range=None, y_range=None, figsize=(8,6), *args, **kwargs):
    """Generate interactive Diane like chart with multiple axes

    Args:
        source (bokeh.ColumnDataSource): Data , e.g downloaded by engine.batch_hist_dataItems(...)
        pltcfg ([list of dicts]): the source columns to plot, and range of y-axis
        x_ax (str, optional): x-axis column as string. Defaults to 'datetime'.
        x_ax_unit (str, optional): unit of x-axis as string. Defaults to None.
        title (str, optional): Main Title of figure. Defaults to None.
        grid (bool, optional): display grid. Defaults to True.
        legend (bool, optional): legend. Defaults to True.
        style (str, optional): style of markers, options i.e. 'line', 'circle'
            circle necessary to enable linked brushing (selection of datapoints)
        x_range (bokeh.figure.x_range, optional): x_range of different bokeh-plot; used to connect x-axis limits
        y_range (bokeh.figure.y_range, optional): y_range of different bokeh-plot; used to connect y-axis limits

    Returns:
        bokeh.plotting.figure: Bokeh plot ready to plot or embed in a layout

    example:
    .....
    from bokeh.io import push_notebook, show, output_notebook
    from bokeh.plotting import figure, output_file, show
    from bokeh.models import LinearAxis, Range1d, HoverTool
    from bokeh.layouts import column, row, gridplot
    from bokeh.models import ColumnDataSource
    from itertools import cycle
    import dmyplant2
    import arrow

    import pandas as pd
    import numpy as np
    import traceback
    import matplotlib
    import sys
    import warnings

    dmyplant2.cred()
    mp = dmyplant2.MyPlant(0)

    # Version mittels Validation Instance
    dval = dmyplant2.Validation.load_def_csv('input.csv')
    vl = dmyplant2.Validation(mp, dval, cui_log=True)
    e = vl.eng_serialNumber(1145166)

    print(f"{e} {e.id}")
    pltcfg=[]
    pltcfg.append( [
        {'col': ['Knock integrator cyl. 07']},
        {'col': ['Ignition voltage cyl. 07']},
        {'col': ['ITP cyl. 07']},
        {'col': ['Exhaust temperature cyl. 07']},
        {'col': ['Operating hours engine']}
    ])

    datastr=[]
    for cfg in pltcfg:
        for y in cfg:
            datastr += y['col']

    ans=datastr_to_dict(datastr)
    dat=ans[0]

    df = mp.hist_data(
        e.id,
        itemIds=dat,
        p_from=arrow.get('2021-03-05 04:00').to('Europe/Vienna'),
        p_to=arrow.get('2021-03-05 05:30').to('Europe/Vienna'),
        timeCycle=1)

    output_notebook()

    df.loc['2021-03-05 05:00':'2021-03-05 06:00']

    title=e._info.get('Validation Engine')

    source = ColumnDataSource(df)
    output_file(title+'.html')
    p=bokeh_chart(source, pltcfg[0], title=title)

    show(p)
    """

    dpi = 66
    mwidth = figsize[0] * dpi
    mheight = figsize[1] * dpi

    #dataitems=pd.read_csv('data/dataitems.csv', sep=';')
    dataitems=dmyplant2.MyPlant.get_dataitems()

    TOOLS = 'pan, box_zoom, xwheel_zoom, box_select, undo, reset, save' #select Tools to display
    colors = cycle(matplotlib.rcParams['axes.prop_cycle']) #colors to use for plot
    linewidth = 2

    if x_ax_unit is not None: #get unit of x_axis either from user or csv-file
        x_unit=x_ax_unit
    else:
        if pd.Series(x_ax).isin(dataitems.myPlantName).any():
            x_unit=dataitems.loc[dataitems.myPlantName==x_ax].iat[0,2]
            if x_unit is np.nan: x_unit=''
        elif x_ax=='Operating hours validation':
            x_unit='h'
        else:
            x_unit=''

    x_axis_label=(f'{x_ax} [{x_unit}]')

    if (x_ax == 'datetime'): #seperate constructors for object for datetime or no datetime x-axis
        p = figure( plot_width=mwidth, plot_height=mheight, x_axis_label=None, x_axis_type='datetime',
        x_range=x_range, y_range=y_range, tools=TOOLS)
    else:
        p = figure( plot_width=mwidth, plot_height=mheight, x_axis_label=x_axis_label,
            tools=TOOLS, x_range=x_range, y_range=y_range)

    if grid==False: p.grid.grid_line_color = None

    p.yaxis.visible = False
    if x_ax=='datetime':
        tooltips = [('Datetime', '@'+x_ax+'{%F %T}')]
    else:
        tooltips = [(x_ax, '@{'+x_ax +'}{0.1 f} '+x_unit)]
    for i, y in enumerate(pltcfg):
        to_remove=[]
        for col in y['col']: #checks if data is available
            #if not pd.Series(col).isin(dataitems.myPlantName).any(): ### instead of comparing with dataitems compare with source
            if col not in source.data: ### instead of comparing with dataitems compare with source
                to_remove.append(col)
                logging.info(f"{col} not found.")
            elif source.data[col].all() == None: #remove of columns if no measurement taken
                to_remove.append(col)
                logging.info(f"{col} not available")
        y['col'] = [e for e in y['col'] if e not in to_remove] #remove elements not contained in dataframe by assigning new list
        if len(y['col'])==0: #jump to next iteration if no col remaining
            continue
        else:
            color = next(cycle(colors))['color']

        if y.get('ylim'):
            ylim = list(y['ylim'])
            p.extra_y_ranges[str(i)] = Range1d(start=ylim[0], end=ylim[1])
        else: #if no ylim defined, use automatic Bokeh Range
            p.extra_y_ranges[str(i)] = DataRange1d()

        unit=[]
        renderers=[]
        for ii, col in enumerate(y['col']):
            if not pd.Series(col).isin(dataitems.myPlantName).any(): #Additional if for handling new data rows generated by function, else is normal behaviour
                if 'unit' in y:
                    unit.append(y['unit'])
                else:
                    unit.append('')
            else:
                unit.append(dataitems.loc[dataitems.myPlantName==col].iat[0,2])

            if pd.isna(unit[-1]):
                unit[-1]=''

            if 'color' in y:
                if isinstance(y['color'], list):
                    color = y['color'][ii]
                else:
                    color = y['color']
            else:
                color = next(cycle(colors))['color']

            # func = getattr(p, style) #to choose between different plotting styles
            # renderers.append(func(source=source, x=x_ax, y=col, #circle or line
            # color=color, y_range_name=str(i), legend_label=col, line_width=linewidth))
            if legend:
                if style == 'line':
                    func = getattr(p, 'line') #to choose between different plotting styles
                    renderers.append(func(source=source, x=x_ax, y=col,  #circle or line
                    color=color, y_range_name=str(i), legend_label=col, line_width=linewidth))
                if style == 'circle':
                    func = getattr(p, 'circle') #to choose between different plotting styles
                    renderers.append(func(source=source, x=x_ax, y=col,  #circle or line
                    color=color, y_range_name=str(i), legend_label=col, line_width=linewidth))
                if style == 'both':
                    func = getattr(p, 'line') #to choose between different plotting styles
                    renderers.append(func(source=source, x=x_ax, y=col,  #circle or line
                    color=color, y_range_name=str(i), legend_label=col, line_width=linewidth))
                    func = getattr(p, 'circle') #to choose between different plotting styles
                    renderers.append(func(source=source, x=x_ax, y=col,  #circle or line
                    color=color, y_range_name=str(i), legend_label=col, line_width=linewidth))
            else:
                if style == 'line':
                    func = getattr(p, 'line') #to choose between different plotting styles
                    renderers.append(func(source=source, x=x_ax, y=col, #circle or line
                    color=color, y_range_name=str(i), line_width=linewidth))
                if style == 'circle':
                    func = getattr(p, 'circle') #to choose between different plotting styles
                    renderers.append(func(source=source, x=x_ax, y=col,  #circle or line
                    color=color, y_range_name=str(i), line_width=linewidth))
                if style == 'both':
                    func = getattr(p, 'line') #to choose between different plotting styles
                    renderers.append(func(source=source, x=x_ax, y=col, #circle or line
                    color=color, y_range_name=str(i), line_width=linewidth))
                    func = getattr(p, 'circle') #to choose between different plotting styles
                    renderers.append(func(source=source, x=x_ax, y=col, #circle or line
                    color=color, y_range_name=str(i), line_width=linewidth))

            tooltips.append((col, '@{'+col +'}{0.2 f} '+unit[-1]))  # or 0.0 a

        if not y.get('ylim'):  #only if y-limits not specified
            p.extra_y_ranges[str(i)].renderers = renderers #only use axis specific renderers for calculation of limits

        if len(unit)==1 or unit.count(unit[0]) == len(unit): #if only one entry or all have the same unit
            llabel = ', '.join(y['col'])+' ['+unit[0]+']'
        else:
            llabel = ', '.join(y['col'])+' ['+', '.join(unit)+']'

        if len(llabel) > 90:
                llabel = llabel[:86] + ' ...'
        if len(y['col']) > 1:
            color = 'black'
        p.add_layout(LinearAxis(y_range_name=str(i),
                            axis_label=llabel, axis_label_text_color=color), 'left')

    callback = CustomJS(code='document.getElementsByClassName("bk-tooltip")[0].style.backgroundColor=“transparent";')
    #callback = CustomJS(code='document.getElementsByClassName("bk-tooltip")[0].style.backgroundColor=“rgba(255,255,255,0.2)";')

    p.add_tools(HoverTool(
        tooltips=tooltips,
        formatters={f'@datetime': 'datetime'}, # use 'datetime' formatter for '@date' field
        mode='mouse',
        callback=callback))  # mode=vline -> display a tooltip whenever the cursor is vertically in line with a glyph

    p.toolbar.active_drag = p.select_one('BoxZoomTool')
    p.toolbar.active_scroll = p.select_one('WheelZoomTool')

    p.legend.click_policy='hide' #hides graph when you click on legend, other option mute (makes them less visible)
    p.legend.location = 'top_left'

    try: #if legend has more than 10 entries reduce spacing
        if len(p.legend.items)>10:
            p.legend.spacing = 0
    except:
        pass

    p.title.text = str(title)
    p.title.text_font_size = '16px'

    return p

def bokeh_chart_engine_comparison(source, pltcfg, variable, eng_names, x_ax='datetime', x_ax_unit=None, title=None, grid=True, legend=True, style='circle', x_range=None, y_range=None, figsize=(8,6), *args, **kwargs):
    """Generate interactive Diane like chart with multiple axes

    Args:
        source (bokeh.ColumnDataSource): Data , e.g downloaded by engine.batch_hist_dataItems(...)
        pltcfg ([list of dicts]): the source columns to plot, and range of y-axis
        variable (String): Variable name to plot
        eng_names (list of strings): Engine names
        x_ax (str, optional): x-axis column as string. Defaults to 'datetime'.
        x_ax_unit (str, optional): unit of x-axis as string. Defaults to None.
        title (str, optional): Main Title of figure. Defaults to None.
        grid (bool, optional): display grid. Defaults to True.
        legend (bool, optional): legend. Defaults to True.
        style (str, optional): style of markers, options i.e. 'line', 'circle'
            circle necessary to enable linked brushing (selection of datapoints)
        x_range (bokeh.figure.x_range, optional): x_range of different bokeh-plot; used to connect x-axis limits
        y_range (bokeh.figure.y_range, optional): y_range of different bokeh-plot; used to connect y-axis limits

    Returns:
        bokeh.plotting.figure: Bokeh plot ready to plot or embed in a layout

    Example:
    pltcfg=[{'col': ['BMW REGENSBURG 5_@_Starts', 'ALPRO M2 616F412 BE_@_Starts', 'BMW REGENSBURG_@_Starts']}]
    Variable='Starts'
    eng_names=['BMW REGENSBURG 5', 'ALPRO M2 616F412 BE', 'BMW REGENSBURG']
    """

    dpi = 80
    mwidth = figsize[0] * dpi
    mheight = figsize[1] * dpi

    dataitems=pd.read_csv('data/dataitems.csv', sep=';')

    TOOLS = 'pan, box_zoom, xwheel_zoom, box_select, undo, reset, save' #select Tools to display
    colors = cycle(matplotlib.rcParams['axes.prop_cycle']) #colors to use for plot
    linewidth = 1
    hovers=[]

    if x_ax_unit is not None: #get unit of x_axis either from user or csv-file
        x_unit=x_ax_unit
    else:
        if pd.Series(x_ax).isin(dataitems.myPlantName).any():
            x_unit=dataitems.loc[dataitems.myPlantName==x_ax].iat[0,2]
            if x_unit is np.nan: x_unit=''
        elif x_ax=='Operating hours validation':
            x_unit='h'
        else:
            x_unit=''

    x_axis_label=(f'{x_ax} [{x_unit}]')

    if (x_ax == 'datetime'): #seperate constructors for object for datetime or no datetime x-axis
        p = figure(
        plot_width=mwidth,
        plot_height=mheight,
        x_axis_label=None,#'datetime',
        x_axis_type='datetime',
        x_range=x_range,
        y_range=y_range,
        tools=TOOLS
        )
    else:
        p = figure(
            plot_width=mwidth,
            plot_height=mheight,
            x_axis_label=x_axis_label,
            tools=TOOLS,
            x_range=x_range,
            y_range=y_range
        )

    if grid==False:
        p.grid.grid_line_color = None

    p.yaxis.visible = False
    tooltips = []
    for i, y in enumerate(pltcfg):
        to_remove=[]
        for col in y['col']: #checks if data is available
            #if not pd.Series(col).isin(dataitems.myPlantName).any(): ### instead of comparing with dataitems compare with source
            if col not in source.data: ### instead of comparing with dataitems compare with source
                to_remove.append(col)
                print (col +' not available! Please check spelling! Not plotted!')
            elif source.data[col].all()==None: #remove of columns if no measurement taken
                to_remove.append(col)
                print (col +' not measured! Can´t be plotted!')
        y['col'] = [e for e in y['col'] if e not in to_remove] #remove elements not contained in dataframe by assigning new list
        if len(y['col'])==0: #jump to next iteration if no col remaining
            continue
        else:
            color = next(cycle(colors))['color']

        if y.get('ylim'):
            ylim = list(y['ylim'])
            p.extra_y_ranges[str(i)] = Range1d(start=ylim[0], end=ylim[1])
        else: #if no ylim defined, use automatic Bokeh Range
            p.extra_y_ranges[str(i)] = DataRange1d()

        unit=[]
        renderers=[]

        for col in y['col']:
            eng_name=col.split('_@_')[0]
            if not pd.Series(variable).isin(dataitems.myPlantName).any(): #Additional if for handling new data rows generated by function, else is normal behaviour
                if 'unit' in y:
                    unit.append(y['unit'])
                else:
                    unit.append('')
            else:
                unit.append(dataitems.loc[dataitems.myPlantName==variable].iat[0,2])

            if unit[-1] is np.nan: unit[-1]=''

            if 'color' in y:
                color = y['color']
            else:
                color = next(cycle(colors))['color']

            func = getattr(p, style) #to choose between different plotting styles
            if style=='circle':
                renderers.append(func(source=source, x=x_ax, y=col, #circle or line
            color=color, y_range_name=str(i), legend_label=eng_name, line_width=linewidth, size=2))
            else:
                renderers.append(func(source=source, x=x_ax, y=col, #circle or line
            color=color, y_range_name=str(i), legend_label=eng_name, line_width=linewidth))
            p.add_tools(HoverTool(tooltips=[(eng_name, '@{'+col +'}{0.2 f} '+unit[-1])], renderers=[renderers[-1]],toggleable=False))

        if not y.get('ylim'):  #only if y-limits not specified
            p.extra_y_ranges[str(i)].renderers = renderers #only use axis specific renderers for calculation of limits

        llabel = variable+' ['+unit[0]+']'

        p.add_layout(LinearAxis(y_range_name=str(i),
                            axis_label=llabel), 'left')

    p.toolbar.active_scroll = p.select_one('WheelZoomTool')

    p.legend.click_policy='hide' #hides graph when you click on legend, other option mute (makes them less visible)
    p.legend.location = 'top_left'

    p.title.text = str(title)
    p.title.text_font_size = '20px'

    return p

def datastr_to_dict (datastr):
    """Generate dict from myPlantNames
    In case name is not valid it gets ignored

    Args:
        datastr (list of str): myPlantNames to be transformed

    Returns:
        dat (dict): dictionary of dataitems
        rename (dict): dict of type {name:myPlantName}

    example:
    .....
    datastr_to_dict(['test123','Exhaust temperature cyl. 23'])

        Output:
        test123 not available! Please check spelling.

        dat={191: ['Exhaust_TempCyl23', 'C (high)']},
        rename={'Exhaust_TempCyl23': 'Exhaust temperature cyl. 23'}"""

    #updated version, can transform myPlantNames from different languages
    data=np.unique(datastr).tolist()

    Request_Ids = pd.read_csv('data/dataitems.csv', sep=';')
    rel_data=pd.DataFrame()

    rename={}
    for da in data:
        # try:
        #     new=dataitems_df.loc[dataitems_df.myPlantName==da]['dataitem'].values[0]
        #     rename [new]=da
        #     da=new
        # except Exception:
        #     pass

        data_id=Request_Ids.loc[Request_Ids['myPlantName']==da]
        if not data_id.empty:
            new=Request_Ids.loc[Request_Ids.myPlantName==da]['name'].values[0]
            rename [new]=da
            rel_data=rel_data.append(data_id)

        #else: #uncommented for less output messages
            #print(da+' not available! Please check spelling.')
            #warnings.warn(da+' not available! Please check spelling.')

    dat = {rec['id']:[rec['name'], rec['unit']] for rec in rel_data.to_dict('records')}
    return dat, rename

def expand_cylinder (y, rel_cyl=all, engi=0):
    """Check if parameter cylinder specific and expand if aplicable

    Args:
        y (dict): one line of a single pltcfg
        rel_cyl (list, optional): Defines relevant cylinders, defaults to all
        engi (dmyplant2.engine, optional): Engine instance to get number of cylinders from

    Returns:
        y (dict): line of a single pltcfg with expanded parameters

    example:
    .....
    """

    if rel_cyl is all:
        if engi != 0:
            e_type=engi.get_property('Engine Type')
            rel_cyl=list(range(1, int(e_type[1:3])+1))
        else:
            rel_cyl=list(range(1, 25))

    add_cyl_short_num=['Inlet valve closure noise', 'Outlet valve closure noise']
    add_cyl_num=['Exhaust temperature','Exhaust temperature delta', 'Ignition voltage', 'ITP','Knock integrator','Knock noise', #'Exhaust temperature delta' added for delta to mean value
    'Pressure 49° before TDC', 'Mechanical noise', 'Cylinder state', 'Close current gradient',
     'Inlet valve closure timing', 'Outlet valve closure timing']
    add_num=['Knock signal','P-max','AI','IMEP','Duration of opening','Conrod bearing temperature','CQ max','CQ','Slow down time']
    add_mid=[]#talk with Sebastian what is looked at analyzis

    to_remove=[]
    for col in y['col']:
        if col in add_cyl_short_num and not col in to_remove:
            for cyl in rel_cyl:
                y['col'].append(f'{col} cyl. {cyl}')
                to_remove.append (col)

        if col in add_cyl_num and not col in to_remove:
            for cyl in rel_cyl:
                y['col'].append(f'{col} cyl. {cyl:02d}')
                to_remove.append (col)

        if col in add_num and not col in to_remove:
            for cyl in rel_cyl:
                y['col'].append(f'{col} {cyl:02d}')
                to_remove.append (col)

        if col in add_mid and not col in to_remove:
            for cyl in rel_cyl:
                y['col'].append(f'{col} cyl. {cyl:02d}')
                to_remove.append (col)

    y['col']=[i for i in y['col'] if not i in to_remove ] #remove original column
    return y

def shrink_cylinder (y, rel_cyl=list(range(1, 25))):
    """Sort out some cylinder specific parameters, so that only the ones interested in are displayed
        The rest is loaded beforehand for shorter overall loading time

    Args:
        y (dict): one line of a single pltcfg
        rel_cyl (list, optional): Defines relevant cylinders, defaults to list:[1,2...,23,24]

    Returns:
        y (dict): line of a single pltcfg with eventually less parameters

    example:
    .....
    """

    rel_cyl=[str(cyl).zfill(2) for cyl in rel_cyl]
    add_cyl_short_num=['Inlet valve closure noise', 'Outlet valve closure noise']
    add_cyl_num=['Exhaust temperature','Exhaust temperature delta', 'Ignition voltage', 'ITP','Knock integrator','Knock noise', #'Exhaust temperature delta' added for delta to mean value
    'Pressure 49° before TDC', 'Mechanical noise', 'Cylinder state', 'Close current gradient',
     'Inlet valve closure timing', 'Outlet valve closure timing']
    add_num=['Knock signal','P-max','AI','IMEP','Duration of opening','Conrod bearing temperature','CQ max','CQ','Slow down time']
    add_mid=[]#talk with Sebastian what is looked at analyzis
    to_check=add_cyl_num+add_num+add_mid

    to_remove=[]
    for col in y['col']:
        if (any(ele in col for ele in to_check) and not col[-2:] in rel_cyl): #check if elemt in expanded elements and not in rel_cyl
            #bug with add_cyl_short_num, exception would need to be added
            to_remove.append (col)

    y['col']=[i for i in y['col'] if not i in to_remove ] #remove original column
    return y

def load_pltcfg_from_excel ():
    """Load plotconfig from Excel Sheet "Input" necessary in same folder

    Returns:
        pltcfg (list of dicts): pltcfg with list of dicts
        plt_titles (list of String): titles of plots
    .....
    """

    import math
    def is_number(s):
        """ Returns True is string is a number. """
        try:
            float(s)
            return math.isfinite(s)
        except ValueError:
            return False

    df_cfg=pd.read_excel('Input_validation_dashboard.xlsx', sheet_name='Pltcfg', usecols=['Plot_Nr', 'Axis_Nr', 'Name', 'Unit', 'y-lim min', 'y-lim max'])
    df_cfg.sort_values(by=['Plot_Nr','Axis_Nr'], inplace=True)
    df_cfg.dropna(subset=['Plot_Nr', 'Axis_Nr', 'Name'], inplace=True)
    df_cfg['p_equal'] = df_cfg.Plot_Nr.eq(df_cfg.Plot_Nr.shift())
    df_cfg['a_equal'] = df_cfg.Axis_Nr.eq(df_cfg.Axis_Nr.shift())

    pltcfg=[]
    plt_titles=[]
    for i in range(len(df_cfg)):
        if df_cfg.p_equal.iloc[i]==False:
            pltcfg.append([]) #new plot
            if df_cfg.Axis_Nr.iloc[i]==0: #append title if axis=0
                plt_titles.append(df_cfg.Name.iloc[i]) #append title
            else:
                plt_titles.append('')

        if df_cfg.Axis_Nr.iloc[i]!=0:
            if df_cfg.a_equal.iloc[i]==False or df_cfg.p_equal.iloc[i]==False:
                pltcfg[-1].append(dict()) #new axis

            y=pltcfg[-1][-1]
            if type(df_cfg.Name.iloc[i])==str:
                if 'col' in y:
                    y['col'].append(df_cfg.Name.iloc[i].replace('\xa0', ' '))
                else:
                    y['col']=[df_cfg.Name.iloc[i].replace('\xa0', ' ')]
                if 'unit' not in y and type(df_cfg.Unit.iloc[i])==str: #take first occurance of unit
                    y['unit']=df_cfg.Unit.iloc[i].replace('\xa0', ' ')

                lim_min=df_cfg['y-lim min'].iloc[i]
                lim_max=df_cfg['y-lim max'].iloc[i]
                if 'ylim' not in y and is_number(lim_min) and is_number(lim_max):
                    y['ylim']=(lim_min, lim_max) #add tuple y lim
    return pltcfg, plt_titles

def show_val_stats (vl, df_loadrange=None, df_starts_oph=None):
    """
    Calculates
    Sort out some cylinder specific parameters, so that only the ones interested in are displayed
    The rest is loaded beforehand for shorter overall loading time

    Args:
        vl (dmyplant2.Validation): Validation Objekt
        df_loadrange (pd.DataFrame) (optional): Dataframe with load information
        df_starts_oph (pd-DatFrame) (optional): DataFrame with information about oph per start

    Returns:
        text_lay (bokeh.models.layouts.Column): Bokeh Column, can be displayed directly with show() or used further in a sheet or tab

    example:
    .....
    """
    from bokeh.models.widgets import DataTable, DateFormatter, TableColumn
    elements=[]
    #### loadrange
    if not df_loadrange.empty:
        loadrange_info=Div(text="<h3>Power load (P/Pnom)</h3>(Valid for displayed data)")

        df_loadrange=df_loadrange*100
        for col in df_loadrange.columns: df_loadrange[col]=df_loadrange[col].map("{:,.1f}%".format)
        df_loadrange.insert(0, 'Engine', df_loadrange.index)
        Columns = [TableColumn(field=Ci, title=Ci) for Ci in df_loadrange.columns] # bokeh columns
        loadrange = DataTable(columns=Columns, source=ColumnDataSource(df_loadrange), autosize_mode='fit_columns', height=30*(len(df_loadrange.index)+1), index_position=None) # bokeh table
        elements+=[loadrange_info, loadrange]#, loadrange_info2]

    #### starts_oph
    if not df_starts_oph.empty:
        starts_oph_info=Div(text="<h3>OPH and Starts</h3>(Valid for displayed data)")

        df_starts_oph['OPH']=df_starts_oph['OPH'].map("{:,.1f}".format)
        df_starts_oph['OPH/ Start']=df_starts_oph['OPH/ Start'].map("{:,.1f}".format)
        df_starts_oph.insert(0, 'Engine', df_starts_oph.index)
        Columns = [TableColumn(field=Ci, title=Ci) for Ci in df_starts_oph.columns] # bokeh columns
        starts_oph = DataTable(columns=Columns, source=ColumnDataSource(df_starts_oph), autosize_mode='fit_columns', height=30*(len(df_starts_oph.index)+1), index_position=None) # bokeh table
        elements+=[starts_oph_info, starts_oph]#, starts_oph_info2]

    d=vl.dashboard
    # Read Values defined in tdef from Myplant into a pd.dataframe
    tdef = {161: 'Count_OpHour', 102: 'Power_PowerAct', 1258: 'OperationalCondition', 19074: 'Various_Bits_CollAlarm'}
    ntable = [[e] + [e.get_dataItem(v) for v in tdef.values()] for e in vl.engines]
    dft = pd.DataFrame(ntable, columns=['Name'] + list(tdef.values()))

    info_text=Div(text="<style>h3, h4{ margin: 0;}</style>"+
    f"<h3>{dft.OperationalCondition.count()} Engines in Validation Fleet:</h3>"+
    f"{dft[((dft.OperationalCondition == 'Running') | (dft.Power_PowerAct > 0))].OperationalCondition.count()} Validation Engines UP and Running<br>"+
    f"{dft[((dft.OperationalCondition != 'Running') & (dft.Power_PowerAct == 0))].OperationalCondition.count()} Validation Engines not Running")

    #Display OPH characteristics
    oph_info_tit=Div(text="<h3>Validation progress</h3>")
    df_oph=pd.DataFrame(columns=['Characteristic','OPH'])
    df_oph=df_oph.append({'Characteristic':'Fleet leader', 'OPH': f"{max(d['OPH Validation']):.0f}"}, ignore_index=True)
    #df_oph=df_oph.append({'Characteristic':'75% quantile', 'OPH': f"{np.quantile(d['oph parts'],q=0.75):.0f}"}, ignore_index=True)
    #df_oph=df_oph.append({'Characteristic':'Median', 'OPH': f"{np.median(d['oph parts']):.0f}"}, ignore_index=True)
    #df_oph=df_oph.append({'Characteristic':'50% quantile', 'OPH': f"{np.quantile(d['oph parts'],q=0.5):.0f}"}, ignore_index=True)
    #df_oph=df_oph.append({'Characteristic':'25% quantile', 'OPH': f"{np.quantile(d['oph parts'],q=0.25):.0f}"}, ignore_index=True)
    df_oph=df_oph.append({'Characteristic':'Average', 'OPH': f"{np.average(d['OPH Validation']):.0f}"}, ignore_index=True)
    df_oph=df_oph.append({'Characteristic':'Cumulated', 'OPH': f"{sum(d['OPH Validation']):.0f}"}, ignore_index=True)

    Columns = [TableColumn(field=Ci, title=Ci) for Ci in df_oph.columns] # bokeh columns
    oph_info = DataTable(columns=Columns, source=ColumnDataSource(df_oph), autosize_mode='fit_columns', height=30*(len(df_oph.index)+1),index_position=None) # bokeh table

    #Displayengines with certain states
    spec_eng=Div(text="<h3>Engines with special states:</h3>")
    tit_run=Div(text="<h4>Engines not running:</h4>")
    dfd=dft[((dft.OperationalCondition != 'Running') & (dft.Power_PowerAct == 0))]
    if dfd.empty:
        table_run=Div(text="<em>All engines running</em>")
    else:
        dfd['Name'] = dfd.apply(lambda row : row[0]._info.get('Validation Engine'), axis=1)
        Columns2 = [TableColumn(field=Ci, title=Ci) for Ci in dfd.columns] # bokeh columns
        table_run = DataTable(columns=Columns2, source=ColumnDataSource(dfd), autosize_mode='fit_columns', height=30*(len(dfd.index)+1)) # bokeh table

    tit_con=Div(text="<h4>Engines without contact:</h4>")
    dfc=dft[((dft.OperationalCondition == 'No Contact') | (dft.OperationalCondition == 'Never Connected'))]

    if dfc.empty:
        table_con=Div(text="<em>All engines in contact</em>")
    else:
        dfc['Name'] = dfc.apply(lambda row : row[0]._info.get('Validation Engine'), axis=1)
        Columns = [TableColumn(field=Ci, title=Ci) for Ci in dfc.columns] # bokeh columns
        table_con = DataTable(columns=Columns, source=ColumnDataSource(dfc), autosize_mode='fit_columns', height=30*(len(dfc.index)+1)) # bokeh table

    tit_alarm=Div(text="<h4>Engines with Alarm Flag not 0:</h4>")
    dfe=dft[dft.Various_Bits_CollAlarm != 0]
    if dfe.empty:
        table_alarm=Div(text="<em>No engines with alarms</em>")
    else:
        dfe['Name'] = dfe.apply(lambda row : row[0]._info.get('Validation Engine'), axis=1)
        Columns3 = [TableColumn(field=Ci, title=Ci) for Ci in dfe.columns] # bokeh columns
        table_alarm = DataTable(columns=Columns3, source=ColumnDataSource(dfe), autosize_mode='fit_columns', height=30*(len(dfe.index)+1)) # bokeh table

    #Display trips
    trip_div=[Div(text="<h3>Recent alarms of engines with alarm:</h3>")]
    dtripped = dft[dft.Various_Bits_CollAlarm == 1]
    for eng in dtripped.values:
        le = eng[0]
        trip_div.append(Div(text='<h4>'+le._info.get('Validation Engine')+'</h4>'))
        dtrips = le.batch_hist_alarms(p_severities=[800], p_offset=0, p_limit=5)
        dtrips['datetime'] = pd.to_datetime(dtrips['timestamp'] * 1000000.0).dt.strftime("%m-%d-%Y %H:%m")
        df_print=dtrips[['datetime', 'message', 'name','severity']]
        Columns = [TableColumn(field=Ci, title=Ci) for Ci in df_print.columns] # bokeh columns
        trip_div.append(DataTable(columns=Columns, source=ColumnDataSource(df_print), autosize_mode='fit_columns', height=30*(len(df_print.index)+1))) # bokeh table

    #Create bar for figure call
    bar_source=ColumnDataSource({'Validation Engines UP and Running': [(dft[((dft.OperationalCondition == 'Running') | (dft.Power_PowerAct > 0))].OperationalCondition.count())], 'Validation Engines not Running': [(dft[((dft.OperationalCondition != 'Running') & (dft.Power_PowerAct == 0))].OperationalCondition.count())]})

    p = figure(plot_width=500,plot_height=50, tools="hover", tooltips="$name: @$name", toolbar_location=None)
    p.axis.visible = False
    p.xgrid.visible = False
    p.ygrid.visible = False
    p.x_range.range_padding = 0
    p.y_range.range_padding = 0
    p.hbar_stack(['Validation Engines UP and Running', 'Validation Engines not Running'], y=10, width=0.9, color=['green', 'grey'], source=bar_source)

    access_time=Div(text='<small>Access time: '+datetime.now().strftime('%d.%m.%y %H:%M')+'</small>')

    elements=elements+[info_text, p, access_time, oph_info_tit, oph_info, spec_eng, tit_run, table_run, tit_con, table_con, tit_alarm, table_alarm, [trip_div]]
    text_lay=layout(children=elements)
    return text_lay

############################# Dieter
def demonstrated_Reliabillity_Plot(vl, beta=1.21, T=30000, s=1000, ft=pd.DataFrame, cl=[10, 50, 90], xmin=None, xmax=None, factor=2.0, ymax=24000):
    """Plot the demonstrated Reliability of the specified validation fleet

    Example:

    ....

    # load input data from files
    dval = pd.read_csv("input2.csv",sep=';', encoding='utf-8')
    dval['val start'] = pd.to_datetime(dval['val start'], format='%d.%m.%Y')
    failures = pd.read_csv("failures.csv",sep=';', encoding='utf-8')
    failures['date'] = pd.to_datetime(failures['date'], format='%d.%m.%Y')

    dmyplant2.demonstrated_Reliabillity_Plot(vl,
            beta=1.21, T=30000, s=1000, ft=failures, cl=[10,50,90], factor=1.3);

    ...

    Args:
        vl ([dmyplant2.Validation class]): [Class with several function around the validation fleet]
        beta (float, optional): [Weibull beta parameter]. Defaults to 1.21.
        T (int, optional): [Runtime for Assessment of Reliabiliy, calculated with LIPSON Method]. Defaults to 30000.
        s (int, optional): [number of points to plot]. Defaults to 1000.
        ft ([type], optional): [pd.DataFrame with observed failures]. Defaults to pd.DataFrame.
            required Columns: date;failures;serialNumber;comment
        cl (list, optional): [list with relialibilty lines for specific confidence levels to plot,
            Numbers between 0 and 100]. Defaults to [10, 50, 90].
        xmin ([timestamp], optional): [left margin of x-axis]. Defaults to None.
        xmax ([timestamp], optional): [right margin of x-axis]. Defaults to None.
        factor (float, optional): [Extrapolation factor]. Defaults to 2.0.
        ymax (int, optional): [right y-axis max value]. Defaults to 24000.

    Raises:
        ValueError: [Time Range not properly specified]
    """
    # define milestones
    start_ts = vl.valstart_ts if xmin == None else xmin  # val start

    # the end of the Plotting interval
    if xmax:
        last_ts = xmax
    else:
        if factor:
            factor = max(factor, 1.0)  # no factor < 1.0 allowed
            elapsed = vl.now_ts - start_ts
            last_ts = start_ts + factor * elapsed
        else:
            raise ValueError("Error in timerange specification.")

    fcol = 'grey'

    # calculate the x axis timerange first
    tr = demonstrated_reliability_sr(vl,
                                     start_ts, last_ts, beta=beta, size=s, ft=ft)[0]  # timestamp x axis start .. end

    # determine the array - index of 'now'
    n_i = _idx(s, start_ts, last_ts, vl.now_ts)

    # create Timerow from Start to 'now'
    n_tr = tr[0:n_i:1]

    # convert to datetime dates - start .. last
    dtr = [datetime.fromtimestamp(t) for t in tr]
    # calculate demonstrated reliability curves for the complete period,
    # confidence intervals CL :
    rel = {c: demonstrated_reliability_sr(vl, start_ts, last_ts,
                                          CL=c/100.0, beta=beta, size=s, ft=ft, T=T)[1] for c in cl}

    # convert to datetime dates - start .. now
    n_dtr = [datetime.fromtimestamp(t) for t in n_tr]
    # copy demontrated reliability values for the validation period up to now:
    n_rel = {c: rel[c][0:n_i:1] for c in cl}

    # define the PLOT
    fig, ax1 = plt.subplots(  # pylint: disable=unused-variable
        figsize=(12, 8), constrained_layout=True)
    # fig, (ax1, ax3) = plt.subplots(2, figsize=(6, 6))

    color = 'tab:red'
    ax1.set_xlabel('date')
    ax1.set_ylabel('Demonstrated Reliability [%]', color=color)
    ax1.set_title('Demonstrated Reliability [%]')

    # now plot the demonstrated reliability curves:
    for CL in cl:
        # complete interval in color fcal
        ax1.plot(dtr, rel[CL], color=fcol, linestyle='-', linewidth=0.5)
        # the current validation interval in multiple colors
        ax1.plot(n_dtr, n_rel[CL], color='red', linestyle='-', linewidth=0.7)

    # define the axis ticks
    ax1.tick_params(axis='y', labelcolor=color)

    # and the axis scales
    ax1.axis((datetime.fromtimestamp(start_ts),
              datetime.fromtimestamp(last_ts), 0, 100))

    # define axis intervals y ...
    ax1.yaxis.set_major_locator(ticker.LinearLocator(13))

    # and x - axis
    locator = dates.AutoDateLocator()
    locator.intervald[dates.MONTHLY] = [1]
    ax1.xaxis.set_major_locator(locator)

    # show a grid
    ax1.grid(color='lightgrey')

    # instantiate a second axes that shares the same x-axis
    ax2 = ax1.twinx()
    ax2.axis((datetime.fromtimestamp(start_ts),
              datetime.fromtimestamp(last_ts), 0, ymax))
    color = 'tab:blue'
    # the x-label was handled with ax1
    ax2.set_ylabel('hours [h]', color=color)
    ax2.tick_params(axis='y', labelcolor=color)
    ax2.yaxis.set_major_locator(ticker.LinearLocator(13))

    # and plot the linearized engine runtime lines vs the 2nd axis
    for e in vl.engines[:]:
        # print(e.Name, e['Engine ID'], e['val start'], e['oph parts'])
        # complete interval in color fcal
        y = [e.oph2(t) for t in tr]
        ax2.plot(dtr, y, linewidth=0.5, color=fcol)
        # the current validation interval in multiple colors
        n_y = [e.oph2(t) for t in n_tr]
        ax2.plot(n_dtr, n_y, label=f"{e.Name} {e['Engine ID']}")

    # NOW plot some Orientation Lines and Test into the Plot

    # possible runtime line
    y = [(t-start_ts) / 3600.0 for t in tr]
    ax2.plot(dtr, y, color='grey', linestyle='--', linewidth=0.7)

    # today line
    ax1.axvline(datetime.now(), color='red',
                linestyle='--', linewidth=0.7)

    # Point of demonstrated reliability at
    # highest Confidence Level, today
    myrel_y = float(
        rel[max(cl)][int((vl.now_ts-start_ts)/(last_ts - start_ts)*s-1)])
    myrel_x = datetime.fromtimestamp(vl.now_ts)
    ax1.scatter(myrel_x, myrel_y, marker='o', color='black', label='point')
    txt = f"CL {max(cl)}%@{T}\nbeta={beta}\nR={myrel_y:.1f}%"

    # some statistical Information.
    myrel_txt_x = datetime.fromtimestamp(vl.now_ts + 200000)
    ax1.text(myrel_txt_x, myrel_y - 9, txt)
    ax1.axis((datetime.fromtimestamp(start_ts),
              datetime.fromtimestamp(last_ts), 0, 120))
    # oph Fleet Leader
    fl = [e.oph2(vl.now_ts) for e in vl.engines]
    fl_point_x = datetime.fromtimestamp(vl.now_ts)
    ax2.scatter(fl_point_x, max(fl), marker='o', color='black', label='point')
    fl_txt_x = datetime.fromtimestamp(vl.now_ts + 200000)
    txt = f'{len(fl)} engines\nmax {max(fl):.0f}h\ncum {sum(fl):.0f}h\navg {statistics.mean(fl):.0f}h\n{arrow.now("Europe/Vienna").format("DD.MM.YYYY HH:mm")}'
    ax2.text(fl_txt_x, max(fl) - T/7, txt)

    # def on_plot_hover(event):
    #     # Iterating over each data member plotted
    #     for curve in ax2.get_lines():
    #         # Searching which data member corresponds to current mouse position
    #         if curve.contains(event)[0]:
    #             print("over %s" % curve.get_gid())

    # plt.legend()
    # fig.canvas.mpl_connect('motion_notify_event', on_plot_hover)

    # TATAAAAH!
    plt.show()

def chart(d, ys, x='datetime', title=None, grid=True, legend=True, notebook=True, *args, **kwargs):
    """Generate Diane like chart with multiple axes

    example:
    .....

    dat = {
        161: ['CountOph','h'],
        102: ['PowerAct','kW'],
        107: ['Various_Values_SpeedAct','U/min'],
        217: ['Hyd_PressCrankCase','mbar'],
        16546: ['Hyd_PressOilDif','bar']
    }

    df = mp.hist_data(
        e.id,
        itemIds=dat,
        p_from=arrow.get('2021-03-05 05:28').to('Europe/Vienna'),
        p_to=arrow.get('2021-03-05 05:30').to('Europe/Vienna'),
        timeCycle=1)

    dmyplant2.chart(df, [
    {'col': ['PowerAct'],'ylim': [0, 5000]},
    {'col': ['Various_Values_SpeedAct'],'ylim': [0, 2500], 'color':'darkblue'},
    {'col': ['CountOph'],'ylim': [0, 500]},
    {'col': ['Hyd_PressCrankCase'],'ylim': [-40, 60]},
    {'col': ['Hyd_PressOilDif'],'ylim': [0, 1]}
    ],
    title = e,
    grid = False,
    figsize = (14,10))

    .....

    Args:
        d (pd.dataFrame): Data , e.g downloaded by engine.batch_hist_dataItems(...)
        ys ([list of dicts]): the DataFrame d columns to plot
        x (str, optional): x-axis column as string. Defaults to 'datetime'.
        title (str, optional): Main Title of figure. Defaults to None.
        grid (bool, optional): displaygrid on left axis. Defaults to True.
        legend (bool, optional): legend. Defaults to True.
    """
    # for entry in kwargs.items():
    #     print("Key: {}, value: {}".format(entry[0], entry[1]))

    fig, ax = plt.subplots(*args, **kwargs)
    fig.patch.set_facecolor('white')

    axes = [ax]
    ax.tick_params(axis='x', labelrotation=30)

    if grid:
        ax.grid()
    if title:
        ax.set_title(title)

    for y in ys[1:]:
        # Twin the x-axis twice to make independent y-axes.
        axes.append(ax.twinx())

    fig.subplots_adjust(top=0.9)
    fig.subplots_adjust(left=0.1)

    extra_ys = len(axes[2:])

    # Make some space on the right side for the extra y-axes.
    if extra_ys > 0:
        if extra_ys > 6:
            print('you are being ridiculous')
            raise ValueError('too many Extra Axes')
        else:
            temp = 0.9 - extra_ys * 0.05

        # print('you are being ridiculous')
        fig.subplots_adjust(right=temp)
        right_additive = 0.065 / temp

        # Move the last y-axis spine over to the right by x% of the width of the axes
        for i, ax in enumerate(axes[2:]):
            ax.spines['right'].set_position(
                ('axes', 1.0 + right_additive * (i+1)))
            ax.set_frame_on(True)
            ax.patch.set_visible(False)
            ax.yaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
        # To make the border of the right-most axis visible, we need to turn the frame
        # on. This hides the other plots, however, so we need to turn its fill off.

    cols = []
    lines = []
    # line_styles = cycle(['-', '-', '-', '--', '-.', ':', 'dotted', ',', 'o', 'v', '^', '<', '>',
    #                     '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd', '|', '_'])
    line_styles = cycle(['-', '-', '-', '--', '-.', ':'])

    colors = cycle(matplotlib.rcParams['axes.prop_cycle'])
    for ax, y in zip(axes, ys):
        ls = next(cycle(line_styles))
        if len(y['col']) == 1:
            col = y['col'][0]
            cols.append(col)
            if 'color' in y:
                    color = y['color']
            else:
                color = next(cycle(colors))['color']
            lines.append(ax.plot(d[x], d[col],
                                 linestyle=ls, label=col, color=color))
            ax.set_ylabel(col, color=color)
            if 'ylim' in y:
                ax.set_ylim(y['ylim'])
            ax.tick_params(axis='y', colors=color)
            ax.spines['right'].set_color(color)
        else:
            for ii, col in enumerate(y['col']):
                if 'color' in y:
                    if isinstance(y['color'], list):
                        color = y['color'][ii]
                    else:
                        color = y['color']
                else:
                    color = next(cycle(colors))['color']
                lines.append(
                    ax.plot(d[x], d[col], linestyle=ls, label=col, color=color))
                cols.append(col)
            llabel = ', '.join(y['col'])
            if len(llabel) > 90:
                llabel = llabel[:87] + ' ..'
            ax.set_ylabel(llabel)
            if 'ylim' in y:
                ax.set_ylim(y['ylim'])
            ax.tick_params(axis='y')
    axes[0].set_xlabel(d.index.name)
    lns = lines[0]
    for l in lines[1:]:
        lns += l
    labs = [l.get_label() for l in lns]
    if legend:
        axes[0].legend(lns, labs, loc=0)
    return fig, ax, axes

# def add_lines(start, lines, ax, *args, **kwargs):
#     ax.axvline(arrow.get(start).shift(seconds=0).datetime, *args, **kwargs)
#     for l in lines:
#         ax.axvline(arrow.get(start).shift(seconds=l).datetime, *args, **kwargs)

def scatter_chart(d, ys, x='datetime', title=None, grid=True, legend=True, notebook=True, *args, **kwargs):
    """Generate Diane like chart with multiple axes

    example:
    .....

    dat = {
        161: ['CountOph','h'],
        102: ['PowerAct','kW'],
        107: ['Various_Values_SpeedAct','U/min'],
        217: ['Hyd_PressCrankCase','mbar'],
        16546: ['Hyd_PressOilDif','bar']
    }

    df = mp.hist_data(
        e.id,
        itemIds=dat,
        p_from=arrow.get('2021-03-05 05:28').to('Europe/Vienna'),
        p_to=arrow.get('2021-03-05 05:30').to('Europe/Vienna'),
        timeCycle=1)

    dmyplant2.chart(df, [
    {'col': ['PowerAct'],'ylim': [0, 5000]},
    {'col': ['Various_Values_SpeedAct'],'ylim': [0, 2500], 'color':'darkblue'},
    {'col': ['CountOph'],'ylim': [0, 500]},
    {'col': ['Hyd_PressCrankCase'],'ylim': [-40, 60]},
    {'col': ['Hyd_PressOilDif'],'ylim': [0, 1]}
    ],
    title = e,
    grid = False,
    figsize = (14,10))

    .....

    Args:
        d (pd.dataFrame): Data , e.g downloaded by engine.batch_hist_dataItems(...)
        ys ([list of dicts]): the DataFrame d columns to plot
        x (str, optional): x-axis column as string. Defaults to 'datetime'.
        title (str, optional): Main Title of figure. Defaults to None.
        grid (bool, optional): displaygrid on left axis. Defaults to True.
        legend (bool, optional): legend. Defaults to True.
    """
    # for entry in kwargs.items():
    #     print("Key: {}, value: {}".format(entry[0], entry[1]))

    fig, ax = plt.subplots(*args, **kwargs)
    fig.patch.set_facecolor('white')

    axes = [ax]
    ax.tick_params(axis='x', labelrotation=30)

    if grid:
        ax.grid()
    if title:
        ax.set_title(title)

    for y in ys[1:]:
        # Twin the x-axis twice to make independent y-axes.
        axes.append(ax.twinx())

    fig.subplots_adjust(top=0.9)
    fig.subplots_adjust(left=0.1)

    extra_ys = len(axes[2:])

    # Make some space on the right side for the extra y-axes.
    if extra_ys > 0:
        if extra_ys > 6:
            print('you are being ridiculous')
            raise ValueError('too many Extra Axes')
        else:
            temp = 0.9 - extra_ys * 0.05

        # print('you are being ridiculous')
        fig.subplots_adjust(right=temp)
        right_additive = 0.065 / temp

        # Move the last y-axis spine over to the right by x% of the width of the axes
        for i, ax in enumerate(axes[2:]):
            ax.spines['right'].set_position(
                ('axes', 1.0 + right_additive * (i+1)))
            ax.set_frame_on(True)
            ax.patch.set_visible(False)
            ax.yaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
        # To make the border of the right-most axis visible, we need to turn the frame
        # on. This hides the other plots, however, so we need to turn its fill off.

    cols = []
    lines = []
    # line_styles = cycle(['-', '-', '-', '--', '-.', ':', 'dotted', ',', 'o', 'v', '^', '<', '>',
    #                     '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd', '|', '_'])
    line_styles = cycle(['-', '-', '-', '--', '-.', ':'])

    colors = cycle(matplotlib.rcParams['axes.prop_cycle'])
    for ax, y in zip(axes, ys):
        ls = next(cycle(line_styles))
        if len(y['col']) == 1:
            col = y['col'][0]
            cols.append(col)
            if 'color' in y:
                color = y['color']
            else:
                color = next(cycle(colors))['color']
            lines.append(
                #ax.plot(d[x], d[col],
                #                 linestyle=ls, label=col, color=color))
                ax.scatter(d[x], d[col], label=col, color=color))
            ax.set_ylabel(col, color=color)
            if 'ylim' in y:
                ax.set_ylim(y['ylim'])
            ax.tick_params(axis='y', colors=color)
            ax.spines['right'].set_color(color)
        else:
            for col in y['col']:
                if 'color' in y:
                    color = y['color']
                else:
                    color = next(cycle(colors))['color']
                lines.append(
                    #ax.plot(d[x], d[col], linestyle=ls, label=col, color=color))
                    ax.scatter(d[x], d[col], label=col, color=color))
                cols.append(col)
            llabel = ', '.join(y['col'])
            if len(llabel) > 90:
                llabel = llabel[:87] + ' ..'
            ax.set_ylabel(llabel)
            if 'ylim' in y:
                ax.set_ylim(y['ylim'])
            ax.tick_params(axis='y')
    axes[0].set_xlabel(d.index.name)
    #lns = lines[0]
    #for l in lines[1:]:
    #    lns = lns + l
    #labs = [l.get_label() for l in lns]
    #if legend:
    #    axes[0].legend(lns, labs, loc=0)

if __name__ == '__main__':
    pass


import numpy as np
from scipy.spatial.transform import Rotation as R
from sympy import symbols
import utility
import cv2
import json

class EPIPOLAR :
    def __init__(self, args):
        self.num_cam = args.num_valid_cam
        self.max_dist_epiline_to_box = args.max_dist_epiline_to_box
        self.max_dist_epiline_cross_to_box = args.max_dist_epiline_cross_to_box
        self.num_valid_cam = args.num_valid_cam

        self.rpn_stride = args.rpn_stride

        self.width = args.width
        self.height = args.height

        self.zoom_out_w = args.rpn_stride * self.width/args.resized_width
        self.zoom_out_h = args.rpn_stride * self.height/args.resized_height

        self.diag = np.sqrt(args.width**2 + args.height**2)

        self.intrin = self.parse_intrin(args.dataset_path) #(num_valid_cam, 3, 3)

        self.args = args

    def parse_intrin(self, path):
        with open(path) as fp:
            j = json.load(fp)

        intrins = j['intrinsics']
        intrins = [np.array(arr).reshape(3, 3) for arr in intrins.values()]
        return np.array(intrins)

    def reset(self, extrins, debug_imgs) :
        self.calc_T_a2b(extrins)
        self.calc_epipole()
        self.debug_imgs = debug_imgs

    def original_pnt_to_resized_pnt(self, pnt):
        x, y = pnt
        x /= self.zoom_out_w
        y /= self.zoom_out_h
        return np.array([x, y])

    def resized_box_to_original_box(self, bbox_list):
        x1 = bbox_list[:, 0] * self.zoom_out_w
        y1 = bbox_list[:, 1] * self.zoom_out_h
        x2 = bbox_list[:, 2] * self.zoom_out_w
        y2 = bbox_list[:, 3] * self.zoom_out_h
        return np.column_stack([x1, y1, x2, y2]).astype('int32')

    def draw_result(self, cam1_idx, cam2_idx, box1, box2, foot, a, b, c):
        src_img = cv2.resize(self.debug_imgs[cam1_idx], (self.width, self.height))
        dst_img = cv2.resize(self.debug_imgs[cam2_idx], (self.width, self.height))
        src_reuslt_img = utility.draw_box(src_img, box1, name = None, color = (0, 0, 255), is_show = False)
        dst_reuslt_img = utility.draw_box(dst_img, box2, name = None, color = (0, 0, 255), is_show = False)
        line_start = (0, int(-c/b))
        line_end = (int(self.width), int(-a*self.width/b - c))
        print('box', box1, 'box', box2, 'line', line_start, line_end)
        dst_reuslt_img = utility.draw_line(dst_reuslt_img, line_start, line_end)
        dst_reuslt_img = cv2.circle(dst_reuslt_img, tuple(map(int, foot)), 5, (0, 0, 255), -1)

        img_list = [src_reuslt_img, dst_reuslt_img]
        concat_img = utility.get_concat_img(img_list, cols=2)
        resized_concat_img = cv2.resize(concat_img, (640, 360))
        cv2.imshow('epipolar', resized_concat_img)
        cv2.waitKey(0)

    def ext_a2b(self, ext_a, ext_b):
        T_a2r = np.eye(4)
        T_a2r[0:3, 0:3] = R.from_euler('xyz', ext_a[3:]).as_dcm()
        T_a2r[0:3, 3] = np.array(ext_a[:3])

        T_b2r = np.eye(4)
        T_b2r[0:3, 0:3] = R.from_euler('xyz', ext_b[3:]).as_dcm()
        T_b2r[0:3, 3] = np.array(ext_b[:3])

        # T_a2b = T_r2b * T_a2r = T_b2r.inv * T_a2r
        T_a2b = np.matmul(np.linalg.inv(T_b2r), T_a2r)

        return T_a2b

    def find_line(self, pt1, pt2):
        x1, y1 = pt1
        x2, y2 = pt2
        d = (y2 - y1) / (x2 - x1)
        e = y1 - x1 * d
        return [-d, 1, -e]

    def find_foot(self, a, b, c, pt):
        x1, y1 = pt
        temp = (-1 * (a * x1 + b * y1 + c) / (a * a + b * b))
        x = temp * a + x1
        y = temp * b + y1
        return [x, y]

    def find_dist(self, pt1, pt2):
        return ((pt1[0] - pt2[0]) ** 2 + (pt1[1] - pt2[1]) ** 2) ** 0.5

    def calc_T_a2b(self, extrin):
        #input : extrin #(num_valid_cam, 6)
        self.T_a2b = np.zeros((self.num_valid_cam, self.num_valid_cam, 4, 4))
        for i in range(self.num_valid_cam):
            for j in range(self.num_valid_cam):
                if(i == j): continue
                self.T_a2b[i, j] = self.ext_a2b(extrin[i], extrin[j])

    def calc_epipole(self) :
        self.epipole = np.zeros((self.num_valid_cam, self.num_valid_cam, 2))
        epipole1_3dpt = np.array([0, 0, 0, 1])
        for i in range(self.num_valid_cam):
            for j in range(self.num_valid_cam) :
                if(i==j) : continue
                epipole1_in2_3dpt = np.matmul(self.T_a2b[i, j], epipole1_3dpt)[:3]
                epipole1_in2_2dpt = np.matmul(self.intrin[j], epipole1_in2_3dpt)
                self.epipole[i, j] = epipole1_in2_2dpt[:2] / epipole1_in2_2dpt[2]

    def get_epipolar_lines(self, ref_cam_idx, ref_box):
        epipolar_line_dict = {}
        for offset in range(1, self.num_cam):
            target_cam_idx = (ref_cam_idx + offset) % self.num_cam
            epipolar_line_dict[target_cam_idx] = self.get_epipolar_line(ref_cam_idx, ref_box, target_cam_idx)

            '''
            a, b, c = self.get_epipolar_line(ref_cam_idx, ref_box, target_cam_idx)
            original_ref_box = self.resized_box_to_original_box(ref_box.reshape(-1, 4))
            ref_box_center = self.get_boxes_centers(original_ref_box)
            foot = self.find_foot(a, b, c, ref_box_center[0])
            self.draw_result(ref_cam_idx, target_cam_idx, original_ref_box[0], original_ref_box[0], foot, a, b, c)
            '''

        return epipolar_line_dict

    def get_epipolar_line(self, ref_cam_idx, ref_box, target_cam_idx):
        ref_T_a2b = self.T_a2b[ref_cam_idx, target_cam_idx]
        ref_epipole_in2_2dpt = self.epipole[ref_cam_idx, target_cam_idx]

        ref_intrin = self.intrin[ref_cam_idx]
        target_intrin = self.intrin[target_cam_idx]

        #print(ref_T_a2b, ref_epipole_in2_2dpt, ref_intrin, target_intrin, ref_box)
        epipolar_line = self.calc_epipolar_line(ref_T_a2b, ref_epipole_in2_2dpt, ref_intrin, target_intrin, ref_box)
        return epipolar_line

    def solve_system_of_equations(self, eq1, eq2):
        a1, b1, c1 = eq1
        a2, b2, c2 = eq2

        A = np.array([[a1, b1], [a2, b2]])
        b = np.array([-c1, -c2])

        x, y = np.linalg.solve(A, b)
        return x, y

    def calc_epipolar_line(self, T_a2b, epipole1_in2_2dpt, cam1_intrin, cam2_intrin, cam1_box):
        cam1_box = self.resized_box_to_original_box(cam1_box.reshape(-1, 4))
        x1, y1, x2, y2 = cam1_box.reshape(4, )
        bbox1_2dpt = (x1 + x2) / 2, (y1 + y2) / 2

        # bbox 1 in camera 2
        bbox1_3dpt = np.matmul(np.linalg.inv(cam1_intrin), np.array([*bbox1_2dpt, 1]))
        bbox1_3dpt = np.array([*bbox1_3dpt.tolist(), 1])

        bbox1_in2_3dpt = np.matmul(T_a2b, bbox1_3dpt)[:3]
        bbox1_in2_2dpt = np.matmul(cam2_intrin, bbox1_in2_3dpt)
        bbox1_in2_2dpt = bbox1_in2_2dpt[:2] / bbox1_in2_2dpt[2]

        # find epipolar line
        a, b, c = self.find_line(bbox1_in2_2dpt, epipole1_in2_2dpt)
        import math
        if math.isnan(a) :
            print('bbox1_2dpt', bbox1_2dpt)
            print('bbox1_in2_3dpt', bbox1_in2_3dpt)
            print('bbox1_in2_2dpt', bbox1_in2_2dpt)

        return a, b, c

    def check_cross_pnt_valid(self, cross_pnt):
        pass

    def get_box_idx_on_cross_line(self, line1, line2, boxes) :
        cross_pnt = self.solve_system_of_equations(line1, line2)
        #if not self.check_cross_pnt_valid(cross_pnt):
        #    return [], [], [-1]
        boxes = self.resized_box_to_original_box(boxes)
        boxes_centers = self.get_boxes_centers(boxes)
        dist = self.find_dist_pnt2pnts(cross_pnt, boxes_centers) / self.diag
        valid_idx = np.where(dist < self.max_dist_epiline_cross_to_box)
        return valid_idx, dist, cross_pnt

    def draw_boxes_with_epiline(self, ref_cam_idx, ref_box, target_cam_idx, epipolar_line, boxes) :
        boxes = self.resized_box_to_original_box(boxes)
        boxes_centers = self.get_boxes_centers(boxes)
        dists = self.find_dist_line2pnts(epipolar_line, boxes_centers)
        dists /= self.diag

        original_ref_box = self.resized_box_to_original_box(ref_box.reshape(-1, 4))[0]

        a, b, c = epipolar_line

        for box, dist in zip(boxes, dists) :
            box_center = ((box[0]+box[2])/2, (box[1]+box[3])/2)
            foot = self.find_foot(a, b, c, box_center)
            print('dist', dist, 'thersh', self.max_dist_epiline_to_box)
            self.draw_result(ref_cam_idx, target_cam_idx, original_ref_box, box, foot, a, b, c)

    def get_box_idx_on_epiline(self, epipolar_line, boxes) :
        boxes = self.resized_box_to_original_box(boxes)
        boxes_centers = self.get_boxes_centers(boxes)
        dist = self.find_dist_line2pnts(epipolar_line, boxes_centers)
        dist /= self.diag
        valid_idx = np.where(dist < self.max_dist_epiline_to_box)
        return valid_idx

    def find_dist_line2pnts(self, line, pnts):
        a, b, c = line
        dist = []
        for pnt in pnts :
            foot = self.find_foot(a, b, c, pnt)
            dist.append(self.find_dist(pnt, foot))
        return np.array(dist)

    def find_dist_pnt2pnts(self, pnt, pnts):
        return np.array( [self.find_dist(pnt, cur_pnt) for cur_pnt in pnts] )

    def get_boxes_centers(self, bboxes):
        x_center = (bboxes[:, 0] + bboxes[:, 2]) / 2
        y_center = (bboxes[:, 1] + bboxes[:, 3]) / 2
        return np.column_stack([x_center, y_center])


#!/usr/bin/env python3

from asist_nsf_2018.experiments import experiments
from asist.wave_probe import read_wave_probe_csv
from asist.utility import running_mean
from asist.pressure import read_pressure_from_netcdf
import matplotlib.pyplot as plt
from matplotlib.dates import num2date, date2num
import numpy as np
import os
from scipy.io import loadmat
from scipy.signal import detrend
from datetime import datetime, timedelta
from dispersion import w2k
from netCDF4 import Dataset
from process_leg import leg1, leg2, leg_slope

plt.rcParams.update({'font.size': 12})

def get_run_elevations(eta, start_index, n, run_length=360, frequency=100, offset=30):
    n0 = start_index + n * run_length * frequency + offset * frequency
    n1 = n0 + run_length * frequency - 2 * offset * frequency
    return eta[n0:n1]

def demean(x):
    return x - np.mean(x)

def variance_spectrum(eta, sampling_rate, fmin=0.1, fmax=100):
    e = demean(detrend(eta))
    n = e.size
    f = np.fft.fftfreq(n, 1 / sampling_rate)[:n//2]
    #df = 2 * sampling_rate / e.size
    df = sampling_rate / e.size
    ai = 2 * np.abs(np.fft.fft(e)[:n//2]) / n
    F = ai**2 / 2 / df
    mask = (f >= fmin) & (f < fmax)
    return F[mask], f[mask], df

def cp_cg(F, f, df, depth):
    w = 2 * np.pi * f
    k = w2k(w, depth)[0]
    cp = w[1:] / k[1:]
    cg = np.diff(w) / np.diff(k)
    return cp, cg

def sig_wave_height(F, df):
    """Significant wave height [m]."""
    return 4 * np.sqrt(np.sum(F * df))

def mean_wave_period(F, f, df):
    """First-order mean wave period [s]."""
    return np.sum(F * df) / np.sum(F * f * df)

def wave_energy(F, df, rhow=1000, g=9.8):
    """Returns total wave energy."""
    return rhow * g * np.sum(F * df)

def radiation_stress(F, f, df, depth, rhow=1000, g=9.8):
    """Returns radiation stress."""
    cp, cg = cp_cg(F, f, df, depth)
    rad_stress_fac = 2 * cg / cp - 0.5
    return rhow * g * np.sum(rad_stress_fac * F[1:] * df)

def mean_water_height(eta, exp, start_index):
    fan, h = [], []
    for n, run in enumerate(exp.runs[:-1]):
        fan.append(run.fan)
        e = get_run_elevations(eta, start_index, n)
        if n == 0:
            offset = np.mean(e)
            h.append(0)
        else:
            h.append(np.mean(e) - offset)
    return np.array(fan), np.array(h)

def mean_slope(h1, h2, dx, rhow=1000, g=9.8, depth=0.42):
    h1, h2 = np.array(h1), np.array(h2)
    hmean = 0.5 * (h1 + h2) + depth
    return rhow * g * hmean * (h2 - h1) / dx

def get_wave_properties(eta, exp, start_index):
    fan, swh, mwp, Sxx = [], [], [], []
    for n, run in enumerate(exp.runs[:-1]):
        e = get_run_elevations(eta, start_index, n)
        F, f, df = variance_spectrum(e, 100)
        fan.append(run.fan)
        swh.append(sig_wave_height(F, df))
        mwp.append(mean_wave_period(F, f, df))
        Sxx.append(radiation_stress(F, f, df, 0.42))
    return np.array(fan), np.array(swh), np.array(mwp), np.array(Sxx)

path = os.environ['WAVEPROBE_DATA_PATH']
L2_DATA_PATH = os.environ['L2_DATA_PATH']
hidrag_path = '/home/milan/Work/sustain/data/hidrag'

frequency = 100 # Hz
run_length = 360 # s

fetch1_c18 = 6.02
fetch2_c18 = 8.71
fetch1_d04 = 4.592
fetch2_d04 = 8.991

dx_c18 = fetch2_c18 - fetch1_c18
dx_d04 = fetch2_d04 - fetch1_d04

# experiments to process
exp_name = 'asist-windonly-fresh'
exp = experiments[exp_name]
known_index = 241500
start_index_fan = 10

_, time, eta3 = read_wave_probe_csv(path + '/' + exp_name + '/ch3.csv')
_, time, eta4 = read_wave_probe_csv(path + '/' + exp_name + '/ch4.csv')
_, time, eta6 = read_wave_probe_csv(path + '/' + exp_name + '/ch6.csv')

start_index = known_index - (start_index_fan // 5)\
                          * run_length * frequency
start_index = 0 if start_index < 0 else start_index

fan, swh3, mwp3, Sxx3 = get_wave_properties(eta3, exp, start_index)
fan, swh4, mwp4, Sxx4 = get_wave_properties(eta4, exp, start_index)
fan, swh6, mwp6, Sxx6 = get_wave_properties(eta6, exp, start_index)

# radiation stress
Sxx3, Sxx4, Sxx6 = map(np.array, [Sxx3, Sxx4, Sxx6])

# radiation stress gradient
dSdx3 = (Sxx6 - Sxx3) / dx_c18
dSdx4 = (Sxx6 - Sxx4) / dx_c18

fan, h3 = mean_water_height(eta3, exp, start_index)
fan, h4 = mean_water_height(eta4, exp, start_index)
fan, h6 = mean_water_height(eta6, exp, start_index)

# mean slopes
s3 = (h6 - h3) / dx_c18
s4 = (h6 - h4) / dx_c18

# air pressure gradient
with Dataset(L2_DATA_PATH + '/air-pressure_asist-christian-shadowgraph.nc') as nc:
    seconds = nc.variables['Time'][:]
    seconds -= seconds[0]
    origin = datetime.strptime(nc.variables['Time'].origin, '%Y-%m-%dT%H:%M:%S')
    time_air = np.array([origin + timedelta(seconds=s) for s in seconds])
    dpdx_air = nc.variables['dpdx'][:]
    fan_air = nc.variables['fan'][:]

exp = experiments['asist-christian-shadowgraph']

dpdx_c18 = []
for run in exp.runs[:-1]:
    t0 = run.start_time + timedelta(seconds=30)
    t1 = run.end_time - timedelta(seconds=30)
    mask = (time_air > t0) & (time_air < t1)
    dpdx_c18.append(np.mean(dpdx_air[mask]))

rhow = 1000
rhoa = 1.15
g = 9.8
depth = 0.42

dpdx_c18 = - np.array(dpdx_c18) / (rhow * g)

### HIDRAG data

# Location of probes from entrance to tank
ps1 = 3.014
ps2 = 7.012
ps3 = 11.009

mat = loadmat(hidrag_path + '/uwvsu2-24.mat')
U = mat['ups'][0]
uw = mat['uw'][0]
LEG3 = mat['LEG3'][0] * 1e-2
LEG1 = mat['LEG1'][0] * 1e-2
M1 = mat['M1'][0]
M3 = mat['M3'][0]

ps13 = - mat['ps13'][0] * 1e-2
ps12 = - mat['ps12'][0] * 1e-2
dpdx = ps13 / (ps3 - ps1)
dpdx -= dpdx[0]

LEG3 -= LEG3[0]
LEG1 -= LEG1[0]

U = np.array([0] + list(U))

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, xlim=(0, 25))
plt.plot(U[1:], LEG1, 'b--', marker='o', ms=5, label='D04, #1, 4.6 m')
plt.plot(U[1:], LEG3, 'b-', marker='o', ms=5, label='D04, #2, 9.0 m')
plt.plot(U, h4, 'r--', marker='o', ms=5, label='C18, #1, 6.0 m')
plt.plot(U, h6, 'r-', marker='o', ms=5, label='C18, #2, 8.7 m')
plt.plot(U, leg1, 'r--', marker='*', ms=10, label='C18, LEG1')
plt.plot(U, leg2, 'r-', marker='*', ms=10, label='C18, LEG2')
plt.plot([0, 50], [0, 0], 'k--')
plt.legend(loc='lower left', fancybox=True, shadow=True)
plt.grid()
plt.xlabel('Wind speed [m/s]')
plt.ylabel('Mean surface elevation [m]')
plt.title('Mean elevation as function of wind speed')
plt.savefig('HIDRAG_elevation.png', dpi=100)
plt.close(fig)

slope_d04 = (LEG3 - LEG1) / dx_d04
slope_c18 = (h6 - h4) / dx_c18

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, xlim=(0, 25))
plt.plot(U[1:], slope_d04, 'b-', marker='o', ms=5, label='D04 dh/dx')
plt.plot(U[1:], dpdx, 'b-', marker='*', ms=8, label='D04 dp/dx')
plt.plot(U, slope_c18, 'r-', marker='o', ms=5, label='C18 dh/dx')
plt.plot(U, leg_slope, 'r-', marker='s', ms=5, label='C18 LEG dh/dx')
plt.plot(U[1:], dpdx_c18[1:], 'r-', marker='*', ms=8, label='C18 dp/dx')
plt.plot([0, 50], [0, 0], 'k--')
plt.legend(loc='upper left', fancybox=True, shadow=True)
plt.grid()
plt.xlabel('Wind speed [m/s]')
plt.ylabel('Elevation and pressure slope')
plt.title('Elevation and pressure slope vs wind speed')
plt.savefig('HIDRAG_slope.png', dpi=100)
plt.close(fig)

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, xlim=(0, 25))
plt.plot(U[1:], M1, 'b--', marker='o', ms=5, label='D04, #1, 4.6 m')
plt.plot(U[1:], M3, 'b-', marker='o', ms=5, label='D04, #2, 9.0 m')
plt.plot(U, Sxx4, 'r--', marker='o', ms=5, label='C18, #1, 6.0 m')
plt.plot(U, Sxx6, 'r-', marker='o', ms=5, label='C18, #2, 8.7 m')
plt.legend(loc='upper left', fancybox=True, shadow=True)
plt.grid()
plt.plot([0, 50], [0, 0], 'k--')
plt.xlabel('Wind speed [m/s]')
plt.ylabel(r'$S_{xx}$ [$kg/s^3$]')
plt.title('Radiation stress $S_{xx}$ vs wind speed')
plt.savefig('HIDRAG_Sxx.png', dpi=100)
plt.close(fig)

dSdx_d04 = (M3 - M1) / dx_d04
dSdx_c18 = (Sxx6 - Sxx4) / dx_c18

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, xlim=(0, 25))
plt.plot(U[1:], dSdx_d04, 'b-', marker='o', ms=5, label='D04')
plt.plot(U, dSdx_c18, 'r-', marker='o', ms=5, label='C18')
plt.legend(loc='upper left', fancybox=True, shadow=True)
plt.grid()
plt.plot([0, 50], [0, 0], 'k--')
plt.xlabel('Wind speed [m/s]')
plt.ylabel(r'$dS_{xx}/dx$ [$N/m^2$]')
plt.title('Radiation stress gradient $dS_{xx}/dx$ vs wind speed')
plt.savefig('HIDRAG_dSdx.png', dpi=100)
plt.close(fig)

# Bottom stress from Brian
taub = rhow * np.array([.0007, .0014, .0013, .0025, .0030, .0038, .0054, .0040, .0061, .01, .0052, 0.0046])**2
taub_c18 = np.array([0] + list(taub))

cd_d04 = (rhow * g * depth * (slope_d04 + dpdx) + dSdx_d04 - taub) / rhoa / (2 * U[1:])**2
cd_c18 = (rhow * g * depth * (slope_c18 + dpdx_c18) + dSdx_c18 - taub_c18) / rhoa / (2 * U)**2
cd_leg = (rhow * g * depth * (leg_slope + dpdx_c18) + dSdx_c18 - taub_c18) / rhoa / (2 * U)**2

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, ylim=(-5e-3, 1.5e-2), xlim=(0, 50))
plt.plot(2 * U[1:], cd_d04, 'b-', marker='o', ms=5, label='D04')
plt.plot(2 * U, cd_c18, 'r-', marker='o', ms=5, label='C18')
plt.plot(2 * U, cd_leg, 'r-', marker='*', ms=10, label='C18 LEG')
plt.legend(loc='lower right', fancybox=True, shadow=True)
plt.grid()
plt.xlabel('Wind speed [m/s]')
plt.ylabel(r'$C_{D10}$')
plt.plot([0, 50], [0, 0], 'k--')
plt.title('Drag coefficient vs wind speed')
plt.savefig('HIDRAG_cd.png', dpi=100)
plt.close(fig)

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, xlim=(0, 25))
plt.plot(U[1:], rhow * g * depth * slope_d04, 'b-', marker='o', ms=5, label='D04, dh/dx')
plt.plot(U[1:], rhow * g * depth * dpdx, 'b-', marker='*', ms=8, label='D04, dp/dx')
plt.plot(U[1:], taub, 'b-', marker='s', ms=8, label=r'D04, $\tau_b$')
plt.plot(U[1:], dSdx_d04, 'b-', marker='v', ms=8, label='D04, dSxx/dx')
plt.plot(U, rhow * g * depth * slope_c18, 'r-', marker='o', ms=5, label='C18, dh/dx')
plt.plot(U, dSdx_c18, 'r-', marker='v', ms=8, label='C18, dSxx/dx')
plt.legend(loc='upper left', fancybox=True, shadow=True)
plt.grid()
plt.xlabel('Wind speed [m/s]')
plt.ylabel('Stress [$N/m^2$]')
plt.plot([0, 50], [0, 0], 'k--')
plt.title('Stress components vs wind speed')
plt.savefig('HIDRAG_mom_budget.png', dpi=100)
plt.close(fig)


import os
# import os.path as pat
import h5py
import numpy as np
import scipy as sp
import pandas as pd
import soundfile as sf
import matplotlib.pyplot as plt

from tqdm import tqdm
from scipy.optimize import minimize, minimize_scalar

from dechorate.dataset import DechorateDataset, SyntheticDataset
from dechorate.cadzow import cadzow_denoise, condat_denoise
from dechorate.utils.dsp_utils import *
from dechorate.utils.file_utils import save_to_pickle, load_from_pickle
from dechorate import constants

dataset_dir = './data/dECHORATE/'
path_to_processed = './data/processed/'
path_to_note_csv = dataset_dir + 'annotations/dECHORATE_database.csv'

datasets = constants['datasets']

def build_all_rirs_matrix_and_annotation(params):

    Fs = params['Fs']
    I = params['I']
    J = params['J']
    D = params['D']
    K = params['K']
    Lt = params['Lt']
    data = params['data']
    L = int(Fs*Lt)

    all_rirs = np.zeros([L, I, J, D])
    toa_sym = np.zeros([K, I, J, D])
    amp_sym = np.zeros_like(toa_sym)
    ord_sym = np.zeros_like(toa_sym)
    wal_sym = np.ndarray([K, I, J, D], dtype=object)
    gen_sym = np.ndarray([K, I, J, D], dtype=object)

    for j in range(0, J):
        for i in tqdm(range(0, I)):
            for d, dataset in enumerate(datasets):

                # get real data
                dset = DechorateDataset(path_to_processed, path_to_note_csv)
                dset.set_dataset(dataset)
                dset.set_entry(i, j)
                mic_pos, src_pos = dset.get_mic_and_src_pos()
                time, h = dset.get_rir()
                assert dset.Fs == Fs

                # get simulation annotation
                synth_dset = SyntheticDataset()
                synth_dset.set_room_size(constants['room_size'])
                synth_dset.set_dataset(dataset)
                synth_dset.set_c(343)
                synth_dset.set_k_order(3)
                synth_dset.set_k_reflc(K)
                synth_dset.set_mic(mic_pos[0], mic_pos[1], mic_pos[2])
                synth_dset.set_src(src_pos[0], src_pos[1], src_pos[2])
                # times, hs = synth_dset.get_rir()
                amp, tau, wall, order, generators = synth_dset.get_note()

                toa_sym[:, i, j, d] = tau
                amp_sym[:, i, j, d] = amp
                wal_sym[:, i, j, d] = wall
                ord_sym[:, i, j, d] = order
                gen_sym[:, i, j, d] = generators

                # plt.plot(np.abs(h[:L])**p + .5*d, label=dataset)
                # if d == 0:
                #     plt.stem(tau*dset.Fs, amp)
                #     for k in range(K):
                #         plt.text(tau[k]*dset.Fs, 0.25, r'$\tau_{%s}^{%d}$' % (wall[k][0], order[k]))

                all_rirs[:, i, j, d] = h[:L]

            # plt.legend()
            # plt.show()

            # ## DIRECT PATH-DECONVOLUTION
            # Fs = dset.Fs
            # rir = rirs[0]
            # dp_idx = np.argmax(rir)
            # L = int(0.2*Fs)
            # rir = rir[:L]
            # def f(x):
            #     x1 = int(x)
            #     dp = rir[dp_idx-x1:dp_idx+x1]
            #     dp_deconv = np.real(np.fft.ifft(np.fft.fft(rir, L) / np.fft.fft(dp, L)))[:L]
            #     cost = np.linalg.norm(dp_deconv**2, ord=1)
            #     return cost
            # res = minimize_scalar(f, bounds=[20, 200], method='Bounded')
            # print(res)
            # print(res.x)

    toa_note = {
        'toa': toa_sym,
        'amp': amp_sym,
        'wall': wal_sym,
        'order': ord_sym,
        'generators' : gen_sym,
    }
    np.save('./data/interim/all_rirs_9srcs.npy', all_rirs)
    save_to_pickle('./data/interim/toa_note_9srcs.pkl', toa_note)
    1/0
    return all_rirs

def direct_path_deconvolution(all_rirs, params):
    L, I, J, D = all_rirs.shape
    all_rirs_dconv = np.zeros_like(all_rirs)
    Fs = params['Fs']

    for i in range(I):
        for j in range(J):
            # assume that the anechoic data have a nice direct path
            anechoic_rir = all_rirs[:, i, j, 0]
            dp_idx = np.argmax(np.abs(anechoic_rir))

            # we consider as direct path everything in the interval peak [-1 ms, +3 ms]
            curr_dp_index = [dp_idx - int(0.001*Fs), dp_idx + int(0.003*Fs)]
            dp = anechoic_rir[curr_dp_index[0]:curr_dp_index[1]]

            for d in range(D):
                rir = all_rirs[:, i, j, d]
                dp_deconv = np.real(np.fft.ifft(np.fft.fft(rir, L) / np.fft.fft(dp, L)))
                # restore the direct path
                offset = int(0.001*Fs)
                # dp_deconv = np.concatenate([np.zeros(offset), dp_deconv])

                all_rirs_dconv[offset:, i, j, d] = dp_deconv[:-offset]
    return all_rirs_dconv

def plot_rir_skyline(rirs, dataset, toa_sym, toa_peak, params):
    K = params['K']
    L, I, J, D = rirs.shape
    IJ = I*J
    # flat the src-mic axis
    c = 0
    rirs_skyline = np.zeros([L, IJ, D])
    toa_sym_skyline = np.zeros([K, IJ, D])
    for j in range(J):
        for i in range(I):
            for d in range(D):
                rir = rirs[:, i, j, d]
                rirs_skyline[:, c, d] = np.abs(rir/np.max(np.abs(rir)))
                toa_sym_skyline[:, c, d] = toa_sym[:, i, j, d]
            c += 1

    for j in range(J):
        plt.axvline(j*params['I'], color='C7')

    # process the Skyline for visualization
    L = 2000
    rirs_skyline = np.clip(rirs_skyline[:L, :, dataset]**2, 0, 0.4)

    plt.imshow(rirs_skyline, extent=[0, I*J, 0, L], aspect='auto')
    for k in range(7):
        plt.scatter(np.arange(IJ)+0.5, L - toa_sym_skyline[k, :, dataset]*params['Fs'], c='C%d'%(k+1), alpha=.6, label='Pyroom DP %d' % k)

    plt.tight_layout()
    plt.legend()
    plt.savefig('./reports/figures/rir_dp_dconv_skyline_after_calibration.pdf')
    plt.show()

def plot_overlapped_rirs(rirs, toa_note, params):
    L, I, J, D = rirs.shape
    K, I, J, D = toa_note['toa'].shape
    Fs = params['Fs']

    for j in range(J):
        for i in range(I):
            plt.figure(figsize=(16, 9))

            for d in range(D):

                rir = rirs[:, i, j, d]

                # if d == 0:
                #     plt.plot(rir**2 + 0.2*d, alpha=.2, color='C1')
                rir_to_plot = (normalize(rir))**2
                rir_to_plot = np.clip(rir_to_plot, 0, 0.34)
                rir_to_plot = normalize(rir_to_plot)
                plt.plot(rir_to_plot + 0.2*d)

                # Print the dataset name
                wall_code_name = 'fcwsen'
                wall_code = [int(i) for i in list(datasets[d])]
                curr_walls = [wall_code_name[w]
                              for w, code in enumerate(wall_code) if code == 1]
                plt.text(50, 0.07 + 0.2*d, datasets[d])
                plt.text(50, 0.03 + 0.2*d, curr_walls)

            # plot the echo information
            for k in range(K):
                toa = toa_note['toa'][k, i, j, d]
                amp = toa_note['amp'][k, i, j, d]
                wall = toa_note['wall'][k, i, j, d]
                order = toa_note['order'][k, i, j, d]
                plt.axvline(x=int(toa*Fs), alpha=0.5)
                plt.text(toa*Fs, 0.025, r'$\tau_{%s}^{%d}$' % (wall.decode(), order), fontsize=12)
            plt.xlim([0, 2000])
            plt.title('RIRs dataset %s\nmic %d, src %d' % (datasets[d], i, j))
            plt.show()

def write_rir_and_note_as_file(rirs, toa_note, params):
    L, I, J, D = rirs.shape
    K, I, J, D = toa_note['toa'].shape
    Fs = params['Fs']

    os.system('mkdir -p ./data/processed/rirs_manual_annotation/')

    for j in range(J):
        print('Processing src', j)

        os.system('mkdir -p ./data/processed/rirs_manual_annotation/src_%d/' % j)

        for i in range(I):

            path = './data/processed/rirs_manual_annotation/src_%d/mic_%d/' % (j,i)
            os.system('mkdir -p ' + path)

            for d in range(D):

                rir = rirs[:, i, j, d]

                rir_to_plot = (normalize(rir))**2
                rir_to_plot = np.clip(rir_to_plot, 0, 0.34)
                rir_to_plot = normalize(rir_to_plot)

                # Print the dataset name
                wall_code_name = 'fcwsen'
                wall_code = [int(i) for i in list(datasets[d])]
                curr_walls = [wall_code_name[w]
                              for w, code in enumerate(wall_code) if code == 1]
                str_curr_wall = '_'.join(curr_walls)
                sf.write(path + '%s_%s_rir.wav' % (datasets[d], str_curr_wall), rir_to_plot, Fs)

            # plot the echo information
            path_to_curr_note = path + '%s_note.txt' % (datasets[d])
            text = ''
            for k in range(K):
                toa = toa_note['toa'][k, i, j, d]
                amp = toa_note['amp'][k, i, j, d]
                wall = toa_note['wall'][k, i, j, d]
                order = toa_note['order'][k, i, j, d]
                text += '%1.6f\t%1.6f\ttau_%s_%d\n' % (toa, toa, wall.decode(), order)

            with open(path_to_curr_note, "w") as text_file:
                text_file.write(text)

# # FFT domain
# nfft = L
# freqs = np.linspace(0, dset.Fs//2, nfft//2)
# freqs_idx = np.arange(0, nfft//2)
# assert len(freqs) == len(freqs_idx)

# fstart, fend = (0, 16000)
# istart = np.argmin(np.abs(freqs-fstart))
# iend = np.argmin(np.abs(freqs-fend))
# print(istart, iend)
# sub_freqs = np.arange(istart, iend, step=1)
# F = len(sub_freqs)
# # frequency_step = np.diff(freqs[sub_freqs])[0]
# frequency_step = (dset.Fs/2)/(L/2-1)
# print('Delta Freq', frequency_step)

# h = np.abs(rirs[R])
# H = np.fft.fft(h, n=nfft)
# plt.plot(np.abs(H))
# H = H[sub_freqs]
# plt.plot(np.abs(np.concatenate([np.zeros(istart), H])))
# plt.show()

# P = F//2
# assert F > 2*K+1

# toep_hs = []
# for i, h in enumerate(rirs):

#     H = np.fft.fft(h)

#     H = H[sub_freqs]

#     if denoising:
#         print('Cadzow Denoising')
#         print(H.shape, P)
#         Th_P = make_toepliz_as_in_mulan(H, P)
#         Th = condat_denoise(Th_P, K, thr_Cadzow=1e-7)
#         # Th = cadzow_denoise(Th, K, thr_Cadzow=1e-7)
#     else:
#         Th = make_toepliz_as_in_mulan(H, K+1)

#     assert Th.shape[1] == K+1
#     toep_hs.append(Th)

# Th = np.concatenate(toep_hs, axis=0).squeeze()
# # Th = toep_hs[0]

# U, Sigma, Vh = np.linalg.svd(Th, full_matrices=False)
# a = np.conj(Vh[-1,:K+1]).squeeze()  # find it in the nullspace

# assert np.allclose(np.linalg.norm(a), 1, atol=1e-3)
# assert len(a) == K+1

# roots = np.roots(a)
# print('Est a')
# print(a.T)

# roots_ref = np.exp(-1j*2*np.pi*frequency_step * tau)
# a_ref = np.poly(roots_ref[::-1])
# print('Ref a')
# print(a_ref.T)

# # print('Annihilation with est', np.linalg.norm(Th @ a.reshape([K+1,1])))
# # print('Annihilation with ref', np.linalg.norm(Th @ a_ref.reshape([K+1, 1])))

# print(1/frequency_step)
# print(Lt)

# tau_est_mod_freq = np.sort(np.mod(np.angle(roots)/(-2*np.pi*frequency_step), 1/frequency_step))

# print('relative peaks ref', tau)
# print('relative peaks mod freq', tau_est_mod_freq)
# print('diff', np.abs(tau - tau_est_mod_freq)*dset.Fs)

# # for rir in rirs:
# # rir = rirs[R]
# for r, rir in enumerate(rirs):
#     # rir = np.abs(rir / np.max(np.abs(rir)))
#     plt.plot(np.arange(len(rir))/dset.Fs, rir**2, label='RIR %d' % r)
# plt.stem(tau, .5*np.ones_like(tau), use_line_collection = True,
#          linefmt='C3-', markerfmt='C3o', label='simulation peaks')
# plt.stem(tau_est_mod_freq, .5*np.ones_like(tau_est_mod_freq), use_line_collection=True,
#          linefmt='C4-', markerfmt='C4x', label='recovered peaks mod freq')
# plt.legend()
# plt.show()

if __name__ == "__main__":
    params = {
        'Fs' : 48000,
        'I' : 30,
        'J' : 9,
        'D' : len(datasets),
        'R' : 0,
        'K' : 25,
        'Lt' : 0.4,
        'data' : ['real', 'synth'][0]
    }

    ## BUILD ALL-RIRs MATRIX
    all_rirs = build_all_rirs_matrix_and_annotation(params)
    1/0
    ## LOAD BACK THE DATA
    all_rirs = np.load('./data/tmp/all_rirs.npy')
    toa_note = load_from_pickle('./data/tmp/toa_note.pkl')

    ## DIRECT-PATH DECOVOLUTION
    all_rirs_clean = direct_path_deconvolution(all_rirs, params)
    np.save('./data/tmp/all_rirs_clean.npy', all_rirs_clean)

    ## RIR SKYLINE
    # plot_rir_skyline(all_rirs_clean, 5, toa_sym, None, params)

    ## PLOT OVERLAPPED RIRS
    plot_overlapped_rirs(all_rirs, toa_note, params)

    ## WRITE RIR as WAV to file for manual annotation with audacity
    # write_rir_and_note_as_file(all_rirs_clean, toa_note, params)

    denoising = True
    concatenate = False
    pass


<reponame>shyam196/egc
import click
import numpy as np
import scipy.sparse as sparse
import aggfuse_cpu

from test import csr_dmm_cpu, csr_fuse_cpu

def random_sparse(n, k, dtype, density, seed=0):
    return sparse.rand(
        n, k, density=density, format="csr", dtype=dtype, random_state=seed
    )

def random_dense(k, n, dtype, seed=0):
    rng = np.random.default_rng(seed)
    return rng.standard_normal(size=(k, n), dtype=dtype)

@click.command()
@click.argument("n", type=int)
@click.argument("k", type=int)
@click.argument("m", type=int)
@click.argument("d", type=float)
def main(n, k, m, d):
    s = random_sparse(n, k, dtype=np.float32, density=d)
    d = random_dense(k, m, dtype=np.float32)
    w = np.zeros((n, 3), dtype=np.float32)
    w[:, 0] = 1

    x_ref = s @ d
    x_reimple = csr_dmm_cpu(s, d)
    x_fuse = csr_fuse_cpu(s, d, w)

    print(x_ref, "\n\n\n\n")
    print(x_reimple, "\n\n\n")
    print(x_fuse)

    print(np.allclose(x_ref, x_reimple))
    diff = x_ref - x_reimple
    print(diff.max(), diff.min())

    print(np.allclose(x_ref, x_fuse))
    diff = x_ref - x_fuse
    print(diff.max(), diff.min())

if __name__ == "__main__":
    main()