[python] plotting misc

import matplotlib.pyplot as plt
import seaborn as sns


SMALL_SIZE = 14
MEDIUM_SIZE = 18
BIGGER_SIZE = 24
BIGGERER_SIZE = 28

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title
plt.rc('axes', titlesize=BIGGER_SIZE)    # fontsize of the figure title

plt.rcParams["figure.figsize"] = [10, 6]  # figsize
plt.rcParams["figure.dpi"] = 100  # default dpi
plt.rcParams["figure.autolayout"] = True
plt.rcParams["legend.loc"] = "upper left"
plt.rcParams["savefig.bbox"] = "tight"
plt.rcParams["figure.autolayout"] = True

plt.rcParams["font.weight"] = "bold"
plt.rcParams["axes.labelweight"] = "bold"


linestyle_str = [
     ('solid', 'solid'),      # Same as (0, ()) or '-'
     ('dotted', 'dotted'),    # Same as (0, (1, 1)) or ':'
     ('dashed', 'dashed'),    # Same as '--'
     ('dashdot', 'dashdot')]  # Same as '-.'

linestyle_tuple = [
     ('dotted',                (0, (1, 1))),
     ('densely dotted',        (0, (1, 1))),

     ('dashed',                (0, (5, 5))),
     ('densely dashed',        (0, (5, 1))),

     ('dashdotted',            (0, (3, 5, 1, 5))),
     ('densely dashdotted',    (0, (3, 1, 1, 1))),

     ('dashdotdotted',         (0, (3, 5, 1, 5, 1, 5))),
     ('densely dashdotdotted', (0, (3, 1, 1, 1, 1, 1)))]



def export_legend(legend, filename="results/legend.pdf", transparent=False):
    fig = legend.figure
    fig.canvas.draw()
    bbox = legend.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(filename, dpi=200, bbox_inches=bbox, transparent=transparent)

ax.tick_params(axis='both', which='major', labelsize=16)

ax.text(0.0, 0.7, f"$r={r}", ha="center", va="center")

fig.text(0.5, 0.04, 'common X', ha='center')
fig.text(0.04, 0.5, 'common Y', va='center', rotation='vertical')


def hue_regplot(data, x, y, hue, palette=None, **kwargs):
    from matplotlib.cm import get_cmap

    regplots = []

    levels = data[hue].unique()

    if palette is None:
        default_colors = get_cmap("tab10")
        palette = {k: default_colors(i) for i, k in enumerate(levels)}

    for key in levels:
        regplots.append(sns.regplot(x=x, y=y, data=data[data[hue] == key], label=key, color=palette[key], **kwargs))

    return regplots


def subplots_centered(nrows, ncols, nfigs, figsize=None, dpi=None):
    """
    Modification of matplotlib plt.subplots(),
    useful when some subplots are empty.

    It returns a grid where the plots
    in the **last** row are centered.

    Inputs
    ------
        nrows, ncols, figsize: same as plt.subplots()
        nfigs: real number of figures
    """
    assert nfigs < nrows * ncols, "No empty subplots, use normal plt.subplots() instead"

    fig = plt.figure(figsize=figsize, dpi=dpi)
    axs = []

    m = nfigs % ncols
    m = range(1, ncols + 1)[-m]  # subdivision of columns
    gs = gridspec.GridSpec(nrows, m * ncols)

    for i in range(0, nfigs):
        row = i // ncols
        col = i % ncols

        if row == nrows - 1:  # center only last row
            off = int(m * (ncols - nfigs % ncols) / 2)
        else:
            off = 0

        ax = plt.subplot(gs[row, m * col + off : m * (col + 1) + off])
        axs.append(ax)

    return fig, np.array(axs)