plotting.py

# standard libraries
import matplotlib.pyplot as plt
from matplotlib import rcParams
from datetime import datetime
import numpy as np
import pandas as pd
import copy
from pathlib import Path
import getpass
from cycler import cycler
from warnings import warn

# external libraries
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.io as pio
try:
    import scienceplots
except ImportError:
    warn('Failed to import "scienceplots".')

# local imports
from .fileoperations import get_relative_path, clean_filename, save_to_pickle
from .stringformatting import params_to_string
from .general import deep_update

# get default plotting style (based on the user name)
USERS_STYLES = {'jlowinski': 'jan', 'fhoffet': 'felix'}
USERNAME = getpass.getuser()
STYLE_DEFAULT = USERS_STYLES.get(USERNAME, 'felix')


def quick_plot(plot_data, plot_with='plotly', style=STYLE_DEFAULT, show=True, sort=True
               , show_title=True, title=None, add_timestamp=True, params=None, data_folder=None, show_grid=None
               , show_legend=None, legend_title=None, legend_sorted=True, legend_location='best'
               , x_label=None, y_label=None, y2_label=None, y2_behind=True
               , append=None
               , **kwargs):
    """
    A function that simplifies plotting with plotly or matplotlib. Suitable only for 2D plots.
    
    Parameters
    ----------
    plot_data : dict or list
        There are several different ways in which `plot_data` can be given:
        - The default is a dict as follows: {'label_1': {'x': ndarray, 'y': ndarray}, 'label_2': {'x': ndarray, 'y': ndarray}}. 
          One can also plot error bars by using keys 'x_err' or 'y_err'. One also add text to data points by using key 'text'.
          Labels will be used as the legend labels.
          One can also add to every data set any keyword argument that is recognized by the chosen plotting library, e.g. {'label_1': {'x': ndarray, 'y': ndarray, 'linewidth': 2, 'linestyle': ':'}}
        - a dict of lists: {'label_1': [x_data, y_data], 'label_2': [x_data, y_data]}
        - a dict of Pandas Series (index is takes as x values): {'label_1': Series, 'label_2': Series}
        In fact one can also mix them freely.
        There are also some simplified ways to pass data, suitable only to plot a single data set:
        - a list: [x_data, y_data] or [y_data]
        - a dict: {'x': ndarray, 'y': ndarray} or {'y': ndarray}
        If one passes only y_data the x_data will be ordinal number of the values in the array.
    plot_with : string, optional
        Can take values 'plotly' or 'matplotlib' (or 'p' and 'm' as their aliases).
    sort : boolean, optional
        If the data is supposed to be sorted based on x_data.
    append : obj or tuple, optional
        One can use this argument to plot extra data in an already existing plot.
        If one wants to use plotly, should pass the figure object. If one wants to use matplotlib, should pass a tuple with figure and all the axes.
        (As it is returned by this function.)
    style : string or typle (string, dict)
        Plotting styles as defined in set_style(). With the dict in the tuple one can override the style settings.
    **kwargs
        Any keyword argument that is recognized by the chosen plotting library.
        
    Returns
    -------
    plotly Figure or tuple of matplotlib (Figure, Axis_1[, Axis_2])
        All the necessary object to modify the figure by hand. Or use it later to append the figure with this function.
        
    """

    # resolve `plot_with` aliases
    plot_with = {'p': 'plotly', 'm': 'matplotlib'}.get(plot_with, plot_with)

    if plot_with not in ['plotly', 'matplotlib']:
        raise ValueError('Incorrect `plot_with`.')
    
    # deal with some info when appending - if nothing was passed, use the original
    if append and plot_with == 'matplotlib':
        try:
            append[0].plot_data
        except AttributeError:  # when appending to a figure that was generated not with quick_plot(), it happens for figures with subplots that have to made by hand with matplotlib
            pass
        #     append_raw = True
        else:
        #     append_raw = False
            title = title if title else append[0].plot_data['extra_data'].get('title_raw', None)
            params = params if params else append[0].plot_data.get('params', None)
            data_folder = data_folder if data_folder else append[0].plot_data.get('data_folder', None)
            legend_title = legend_title if legend_title else append[0].plot_data.get('legend_title', None)
            x_label = x_label if x_label else append[0].plot_data.get('x_label', None)
            y_label = y_label if y_label else append[0].plot_data.get('y_label', None)
    
    # keep all the parameters that were used to create the plot and later save with the figure
    d = locals()
    d.pop('append')  # `append` might be a class instance and we don't want to copy it
    # d.pop('append_raw', None)  # flag that should not be stored
    dict_to_save = copy.deepcopy(d)
    # to have kwargs not as a dictionary but rather as elements of `dict_to_save`
    dict_to_save.pop('kwargs', None)
    if kwargs:
        # legacy thing - some named arguments that are not supported any more
        kwargs.pop('return_figure', None)
        kwargs.pop('title_raw', None)
        # if one appends, extra_data causes problems
        kwargs.pop('extra_data', None)
        
        dict_to_save.update(kwargs)
    
    # to keep some extra data    
    dict_to_save['extra_data'] = {}

    plot_data = copy.deepcopy(plot_data)  # to be sure that any operation done on the data des not affect the original data
    
    # set the plotting style
    if isinstance(style, (tuple, list)):  # to deal with the case when `style` is a tuple
        set_style(style[0], plot_with, **style[1])
        style = style[0]
    else:
        set_style(style, plot_with)
    
    # prepare the data
    
    # one can provide `plot_data` as list - convert it into a dictionary
    if isinstance(plot_data, list) or isinstance(plot_data, tuple):
        if len(plot_data) == 1:
            plot_data = {'y': plot_data[0]}
        elif len(plot_data) == 2:
            plot_data = {'x': plot_data[0], 'y': plot_data[1]}
            
    # one can also provide a dictionary of lists - convert it to match the expected data structure
    for k, v in plot_data.items():
        if (isinstance(v, list) or isinstance(v, tuple)) and (k != 'x' and k != 'y'):
            plot_data[k] = {'x': v[0], 'y': v[1]}
                
    # check if there only single data set without a label - add a dummy label; also set the visibility of the legend
    if 'y' in plot_data:
        # add simple 'x', if none is provided
        if 'x' not in plot_data: 
            plot_data['x'] = np.arange(np.size(plot_data['y']))
        plot_data = {'label': plot_data}

    # make sure that everything is an numpy array
    for label in plot_data:
        if 's' in plot_data[label]:
            series = plot_data[label].pop('s')
            plot_data[label]['x'] = series.index.to_numpy(dtype=float)
            plot_data[label]['y'] = series.to_numpy(dtype=float)
        elif isinstance(plot_data[label], pd.Series) :
            series = plot_data.pop(label)
            plot_data[label] = {}
            plot_data[label]['x'] = series.index.to_numpy(dtype=float)
            plot_data[label]['y'] = series.to_numpy(dtype=float)
        else:
            plot_data[label]['x'] = np.array(plot_data[label]['x'])
            plot_data[label]['y'] = np.array(plot_data[label]['y'])

        if sort:
            # sort x
            index_array = np.argsort(plot_data[label]['x'])
            plot_data[label]['x'] = plot_data[label]['x'][index_array]
            # sort y accordingly
            plot_data[label]['y'] = plot_data[label]['y'][index_array]
            for key in ['y_err', 'x_err', 'text']:  # check for optional keys
                if key in plot_data[label]:
                    plot_data[label][key] = np.array(plot_data[label][key])
                    plot_data[label][key] = plot_data[label][key][index_array]

    # update the data dictionary (so it would be always in a consistent format)
    dict_to_save['plot_data'] = plot_data

    # simple package specific settings
    if plot_with == 'plotly':
        new_line = '<br>'
    elif plot_with == 'matplotlib':
        new_line = '\n'

    # generate or modify the plot title
    if show_title:
        
        # keep the original title
        dict_to_save['extra_data']['title_raw'] = title

        if title is None:
            title = ''

        if data_folder is not None:
            if title: title += new_line
            relative_folder = get_relative_path(data_folder, depth=3)
            title += f'source: {relative_folder}'

        if add_timestamp:
            if title: title += new_line
            plotting_timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            title += f'plotted: {plotting_timestamp}'

        if params:
            if title: title += new_line
            title += f'parameters: {params_to_string(params)}'
            
        # update the title
        dict_to_save['title'] = title
                    
    basic_keys = ['x', 'y', 'x_err', 'y_err', 'second_y', 'text', 's']  # dict keys that do NOT specify the layout
    
    # get the default parameters controlling the layout
    layout_dict = get_layout_dict(style, plot_with, plot_data)
    # override `layout_dict` with the general kwargs and then with kwargs for a given data series
    for label in layout_dict:
        layout_dict[label].update(kwargs)
        layout_dict[label].update({k: v for (k, v) in plot_data[label].items() if k not in basic_keys})

    if plot_with == 'plotly':
        # fig = go.Figure()
        if append is None:
            fig = make_subplots(specs=[[{"secondary_y": True}]])
        else:
            fig = append
            
        # axes = {'primary': None}  # to keep the return structure consistent with matplotlib case
        
        for label in plot_data:

            # set the plotting mode
            mode = 'lines' if plot_data[label]['x'].size > 200 else 'markers+lines'
            if 'text' in plot_data[label]:
                mode += '+text' 
            
            plot_dict = dict(x=plot_data[label]['x']
                             , y=plot_data[label]['y']
                             , name=label
                             , mode=mode
                             , legendgroup=label
                             )
            
            # # general layout parameters
            # plot_dict.update(kwargs)

            # # layout parameters of a give data series
            # layout_dict = {}
            # for key in plot_data[label]:
            #     if key not in basic_keys:
            #         layout_dict[key] = plot_data[label][key]
            
            # apply layout settings        
            plot_dict.update(layout_dict[label])
            
            # add optional data
            if 'y_err' in plot_data[label]:
                plot_dict.update({'error_y': dict(type='data', visible=True, array=plot_data[label]['y_err'])})
            if 'x_err' in plot_data[label]:
                plot_dict.update({'error_x': dict(type='data', visible=True, array=plot_data[label]['x_err'])})
            if 'text' in plot_data[label]:
                plot_dict.update({'text': plot_data[label]['text']})
                
            # choose to which axis add the trace
            if plot_data[label].get('second_y', False):
                fig.add_trace(go.Scatter(**plot_dict), secondary_y=True)
                # axes['secondary'] = None  # to keep the return structure consistent
            else:
                fig.add_trace(go.Scatter(**plot_dict), secondary_y=False)            

        if show_title:
            fig.update_layout(title=title)

        if legend_title:
            fig.update_layout(legend_title_text=legend_title)

        if x_label:
            fig.update_layout(xaxis_title=x_label)

        if y_label:
            fig.update_layout(yaxis_title=y_label)

            
    elif plot_with == 'matplotlib':
        # get fig and axis
        if append is None:  # new ones
            fig, ax = plt.subplots(1, 1)
            axes = {'primary': ax}
            count = 0  # how many traces are already in the plot (new plot - zero traces)
        else:  # from figure that is being appended
            if len(append) == 2:
                fig, ax = append
                axes = {'primary': ax}
            elif len(append) == 3:
                fig, ax, ax2 = append
                axes = {'primary': ax, 'secondary': ax2}
                
            # try to add the new plot data to the already existing one
            try:
                fig.plot_data
            except AttributeError:
                pass
            else:
                # check if the all the labels are distinct (if not append some random integer to them)
                # GIVES ERRORS !!!
                for k in plot_data:
                    if k in fig.plot_data['plot_data']:
                        random_int = int(np.random.random() * 1000000)
                        plot_data[f'{k} {random_int}'] = plot_data.pop(k)
                
                count = len(fig.plot_data['plot_data'])  # how many traces are already in the plot
                fig.plot_data['plot_data'].update(plot_data)
                
        # # extra (conditional) styling
        # layout_dict = {label: {} for label in plot_data}
        # if style in ['jan', 'jan_thesis']:
        #     for label in plot_data:
        #         if plot_data[label]['x'].size >= 100:
        #             layout_dict[label]['marker'] = ''
        #             layout_dict[label]['linestyle'] = '-'
        #         else:
        #             if plot_data[label].get('second_y', False):
        #                 layout_dict[label]['marker'] = 'x'
        #                 layout_dict[label]['linestyle'] = ':'
        #             else:
        #                 layout_dict[label]['marker'] = 'o'
        #                 layout_dict[label]['linestyle'] = '--'
                        
        # # override `layout_dict` with the general kwargs and then with kwargs for a given data series
        # for label in layout_dict:
        #     layout_dict[label].update(kwargs)
        #     layout_dict[label].update({k: v for (k, v) in plot_data[label].items() if k not in basic_keys})
                    
        # plot data series
        for j, label in enumerate(plot_data):
                                    
            # choose y axis
            if plot_data[label].get('second_y', False):
                if 'secondary' in axes:
                    axis = axes['secondary']
                else:  # if the secondary axis doesn't exist, create it
                    axis = ax.twinx()
                    axes['secondary'] = axis
            else:
                axis = ax

            # dict that keeps all the details of a given trace 
            plot_dict = {'label': label}
            
            # add layout info
            plot_dict.update(layout_dict[label])

            # add error bars
            if 'x_err' in plot_data[label]:
                plot_dict.update({'xerr': plot_data[label]['x_err']})
            if 'y_err' in plot_data[label]:
                plot_dict.update({'yerr': plot_data[label]['y_err']})

            # plot
            if 'y_err' in plot_data[label] or 'x_err' in plot_data[label]:
                axis.errorbar(plot_data[label]['x'], plot_data[label]['y'], **plot_dict)
            else:
                axis.plot(plot_data[label]['x'], plot_data[label]['y'], **plot_dict)

            # add annotations
            if 'text' in plot_data[label]:
                for i, txt in enumerate(plot_data[label]['text']):
                    axis.annotate(txt, (plot_data[label]['x'][i], plot_data[label]['y'][i]))
                                
        if show_grid is None:
            axes['primary'].grid(True)
        else:
            axes['primary'].grid(show_grid)

        if show_title:
            axes['primary'].set_title(title, loc='left')
        
        # if show_legend:
        if 'secondary' in axes:
            lines1, labels1 = axes['primary'].get_legend_handles_labels()
            lines2, labels2 = axes['secondary'].get_legend_handles_labels()
            lines, labels = lines1 + lines2, labels1 + labels2
        else:
            lines, labels = axes['primary'].get_legend_handles_labels()

        if legend_sorted:
            labels, lines = zip(*sorted(zip(labels, lines)))

        axes['primary'].legend(lines, labels, frameon=False, title=legend_title, loc=legend_location)

        if x_label:
            axes['primary'].set_xlabel(x_label)

        if y_label:
            axes['primary'].set_ylabel(y_label)

        if 'secondary' in axes and y2_label:
            axes['secondary'].set_ylabel(y2_label)

        # makes the traces of the primary axis appear on top of the traces of the secondary axis
        if 'secondary' in axes and y2_behind:
            axes['primary'].set_zorder(1)  # default zorder is 0 for ax1 and ax2
            axes['primary'].set_frame_on(False)  # prevents ax1 from hiding ax2

        # change background of the plot to white (instead of the default transparent)
        fig.patch.set_facecolor('white')
        
    # decide if the legend should be shown          
    if len(plot_data) == 1 and 'label' in plot_data:
        show_legend = False if show_legend is None else show_legend
    else:
        show_legend = True if show_legend is None else show_legend
    if plot_with == 'plotly':
        if not show_legend:
            fig.update_layout(showlegend=False)
    elif plot_with == 'matplotlib':
        if not show_legend:
            axes['primary'].get_legend().remove()
        
    # show, return, save
           
    if show and plot_with == 'plotly':
        fig.show()

    # save the plot data (works only for matplotlib)
    if plot_with == 'matplotlib':
        if append:  # and not append_raw:
            try:
                fig.plot_data
            except AttributeError:  # handle the case when the figure was not created with quick_plot()
                fig.plot_data = dict_to_save
            else:
                fig.plot_data = deep_update(fig.plot_data, dict_to_save)
        else:
            fig.plot_data = dict_to_save
            
        fig.save_figure = save_figure_matplotlib.__get__(fig)

        # alias
        fig.save = fig.save_figure

        if x_label and y_label:
            fig.filename_default = f'{y_label}_vs_{x_label}'
        
    # return the figure object
    if plot_with == 'plotly':
        return fig
    else:
        if 'secondary' in axes:
            return fig, axes['primary'], axes['secondary']  
        else:
            return fig, axes['primary']


# def save_figure_plotly(fig, filename, save_data=True):

#     if save_data:
#         save_to_pickle(fig.plot_data, filename, subfolder='plot_data')

#     filename = str(filename)
#     fig.write_image(filename)


def save_figure_matplotlib(fig, filename=None, save_data=True, override=False, extra_name=None, save_folder=None, format='png'):

    if filename is None:
        try:
            filename = fig.filename_default
        except AttributeError:
            raise ValueError('Missing `filename` to save the figure.')

    if save_folder:
        filename = Path(save_folder) / filename

    filename = clean_filename(filename)

    if extra_name:
        filename = Path(filename.stem + f'-{extra_name}' + filename.suffix)

    if filename.suffix == '':
        filename = filename.with_suffix(f'.{format}')

    if Path(filename).is_file() and not override:
        raise ValueError(f'Failed to save results to a file, because {filename.name} already exists. If one wants to override, should set override=True.')

    fig.savefig(filename)

    if save_data:
        save_to_pickle(fig.plot_data, filename, subfolder='plot_data')


# alias
def plot(*args, **kwargs):
    return quick_plot(*args, **kwargs)


def set_style(style, plot_with, **kwargs):
    
    # matplotlib styles
    
    # reset first to the default style
    if style != 'no_reset':
        plt.style.use('default')
    
    # choose style which will be modified
    if style == 'jan_thesis':
        plt.style.use(['science', 'grid', 'bright'])  # from 'scienceplots'
    
    styles_matplotlib = {'default': {}
                         , 'no_reset': {}
                         , 'jan': {'font.family': 'serif'
                                   , 'font.size': 10
                                   , 'figure.figsize': (9., 9. * (np.sqrt(5.0) - 1.0) / 2.0)
                                   , 'figure.dpi': 300
                                   , 'figure.autolayout': True
                                   , 'mathtext.fontset': 'cm'
                                   , 'axes.titlesize': 10
                                   , 'axes.labelsize': 18
                                   , 'axes.prop_cycle': cycler(color=['#e31a1c', '#ff7f00', '#cab2d6', '#6a3d9a', '#eded62', '#b15928', '#a6cee3', '#1f78b4', '#b2df8a', '#33a02c', '#fb9a99', '#fdbf6f'])
                                   , 'legend.fontsize': 10
                                   , 'legend.title_fontsize': 10
                                   , 'xtick.labelsize': 18
                                   , 'ytick.labelsize': 18
                                   , 'lines.markersize': 5
                                   , 'lines.markeredgewidth': 2
                                   , 'markers.fillstyle': 'full'
                                   , 'savefig.pad_inches': 0.1
                                   , 'savefig.bbox': 'tight'
                                   }
                         , 'felix': {'text.usetex': True
                                   , 'font.family': 'serif'
                                   , 'figure.figsize': (3 + 3/8, (3 + 3/8) * (np.sqrt(5.0) - 1.0) / 2.0)
                                   , 'figure.dpi': 300
                                   , 'axes.prop_cycle': cycler(color=['#E97451', '#EBB908', '#50E1AD', '#6a3d9a', '#eded62', '#b15928', '#a6cee3', '#1f78b4', '#b2df8a', '#33a02c', '#fb9a99', '#fdbf6f'],
                                                               mfc = ['#F6BAA8', '#F5E7B6', '#AFF6DD', '#6a3d9a', '#eded62', '#b15928', '#a6cee3', '#1f78b4', '#b2df8a', '#33a02c', '#fb9a99', '#fdbf6f'])
                                   , 'legend.fontsize': 6
                                   , 'legend.handlelength': 1.5
                                   , 'legend.labelspacing': 0.2
                                   , 'ytick.direction': 'in'
                                   , 'xtick.direction': 'in'
                                   , 'ytick.minor.visible': True
                                   , 'xtick.minor.visible': True
                                   , 'xtick.top': True
                                   , 'ytick.right': True
                                   , 'savefig.pad_inches': 0.03
                                   , 'savefig.bbox': 'tight'
                                   , 'xtick.labelsize': 8
                                   , 'ytick.labelsize': 8
                                   , 'axes.labelsize': 9
                                   , 'ytick.major.pad': 1.2
                                   , 'xtick.major.pad': 1.2
                                   , 'axes.labelpad': 1.5
                                   , 'axes.spines.right': True
                                   , 'axes.spines.left': True
                                   , 'axes.spines.top': True
                                   , 'axes.spines.bottom': True
                                   , 'axes.titlesize': 7
                                   , 'font.size': 7
                                   , 'markers.fillstyle': 'full'
                                   , 'lines.markersize': 4
                                   , 'lines.markeredgewidth': 0.7
                                   , 'lines.linewidth': 1
                                   , 'grid.linewidth': 0.3
                                   , 'axes.linewidth': 1
                                   , 'figure.autolayout' : True}
                         , 'felix_subplots': {'text.usetex': True
                                   , 'font.family': 'serif'
                                   , 'figure.figsize': (3, 1.7)
                                   , 'figure.dpi': 250
                                   , 'axes.prop_cycle': cycler(color=['#E97451', '#EBB908', '#50E1AD', '#6a3d9a', '#eded62', '#b15928', '#a6cee3', '#1f78b4', '#b2df8a', '#33a02c', '#fb9a99', '#fdbf6f'],
                                                               mfc = ['#F6BAA8', '#F5E7B6', '#AFF6DD', '#6a3d9a', '#eded62', '#b15928', '#a6cee3', '#1f78b4', '#b2df8a', '#33a02c', '#fb9a99', '#fdbf6f'])
                                   , 'legend.fontsize': 5
                                   , 'legend.handlelength': 1.5
                                   , 'legend.labelspacing': 0.2
                                   , 'ytick.direction': 'in'
                                   , 'xtick.direction': 'in'
                                   , 'ytick.minor.visible': True
                                   , 'xtick.minor.visible': True
                                   , 'xtick.top': True
                                   , 'ytick.right': True
                                   , 'savefig.pad_inches': 0.03
                                   , 'savefig.bbox': 'tight'
                                   , 'xtick.labelsize': 8
                                   , 'ytick.labelsize': 8
                                   , 'axes.labelsize': 9
                                   , 'ytick.major.pad': 1.2
                                   , 'xtick.major.pad': 1.2
                                   , 'axes.labelpad': 1.5
                                   , 'axes.spines.right': True
                                   , 'axes.spines.left': True
                                   , 'axes.spines.top': True
                                   , 'axes.spines.bottom': True
                                   , 'axes.titlesize': 7
                                   , 'font.size': 7
                                   , 'markers.fillstyle': 'full'
                                   , 'lines.markersize': 4
                                   , 'lines.markeredgewidth': 0.7
                                   , 'lines.linewidth': 1
                                   , 'grid.linewidth': 0.3
                                   , 'axes.linewidth': 1
                                   , 'figure.autolayout' : True}
                         , 'jan_thesis': {'figure.figsize': (4.927 / 2., 4.927 / 2. * (np.sqrt(5.0) - 1.0) / 2.0)
                                          , 'markers.fillstyle': 'none'
                                          , 'font.size': 10
                                        #   , 'xtick.labelsize': 9
                                        #   , 'ytick.labelsize': 9
                                          , 'lines.markersize': 3
                                          , 'lines.markeredgewidth': 0.7
                                          , 'savefig.pad_inches': 0.02
                                          , 'savefig.bbox': 'tight'
                                          , 'axes.labelpad': 2.
                                          , 'axes.grid': False
                                          , 'xtick.major.pad': 2.8
                                          , 'xtick.minor.pad': 2.8
                                          , 'ytick.major.pad': 1.8
                                          , 'ytick.minor.pad': 1.8
                                          , 'legend.handlelength': 1.
                                          , 'legend.labelspacing': 0.1
                                          , 'legend.handletextpad': 0.4
                                          }
                        , 'lukas': {'font.family': 'serif'
                                  , 'font.size': 8
                                  , 'figure.figsize': (2,1)
                                  , 'figure.dpi': 300
                                  , 'mathtext.fontset': 'cm'
                                  , 'axes.titlesize': 8
                                  , 'axes.labelsize': 8
                                  , 'axes.prop_cycle': cycler(color=['#e31a1c', '#ff7f00', '#cab2d6', '#6a3d9a', '#eded62', '#b15928', '#a6cee3', '#1f78b4', '#b2df8a', '#33a02c', '#fb9a99', '#fdbf6f'])
                                  , 'legend.fontsize': 8
                                  , 'legend.title_fontsize': 8
                                  , 'xtick.labelsize': 8
                                  , 'ytick.labelsize': 8
                                  , 'lines.markersize': 5
                                  , 'lines.markeredgewidth': 2
                                  , 'markers.fillstyle': 'full'
                                  , 'savefig.pad_inches': 0.1
                                  , 'savefig.bbox': 'tight'
                                  }
                         }
    
    # pyplot styles - to define you own plotly style read https://plotly.com/python/templates/
    styles_plotly = {}
    
    # select the style
    if plot_with == 'plotly':
        style = styles_plotly.get(style, 'plotly')
    elif plot_with == 'matplotlib':
        style = styles_matplotlib.get(style, styles_matplotlib['default'])
        # update the selected style with kwargs
        style.update(kwargs)

    # set the style
    if plot_with == 'plotly':
        pio.templates.default = style
    elif plot_with == 'matplotlib':
        rcParams.update(style)


def subplots(nrows=1, ncols=1, style=STYLE_DEFAULT, **kwargs):
    """
    It modifies matplolib function subplots() to make it suitable for appending with quick_plot().
    It works same as the original function. There is one extra argument `style` which works identically as in quick_plot().
    """
    
    # set the plotting style
    if isinstance(style, (tuple, list)):  # to deal with the case when `style` is a tuple
        set_style(style[0], 'matplotlib', **style[1])
    else:
        set_style(style, 'matplotlib')
    
    fig, axes = plt.subplots(nrows, ncols, **kwargs)
    
    # fig.plot_data = {'extra_data': {}}  # to match the expected data structure
    
    return fig, axes


def get_layout_dict(style, plot_with, plot_data):
    
    layout_dict = {label: {} for label in plot_data}
    
    if plot_with == 'plotly':
        pass
    
    elif plot_with == 'matplotlib':        
        if style in ['jan', 'jan_thesis']:
            for label in plot_data:
                if plot_data[label]['x'].size >= 100:
                    layout_dict[label]['marker'] = ''
                    if plot_data[label].get('second_y', False):
                        layout_dict[label]['linestyle'] = '--'
                    else:
                        layout_dict[label]['linestyle'] = '-'
                else:
                    if plot_data[label].get('second_y', False):
                        layout_dict[label]['marker'] = 'x'
                        layout_dict[label]['linestyle'] = ':'
                    else:
                        layout_dict[label]['marker'] = 'o'
                        layout_dict[label]['linestyle'] = '--'
                        
    return layout_dict