"""A set of convenience functions used for producing plots in `dabest`."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/API/plot_tools.ipynb.

# %% ../nbs/API/plot_tools.ipynb 2
from __future__ import annotations

# %% auto 0
__all__ = ['halfviolin', 'get_swarm_spans', 'error_bar', 'check_data_matches_labels', 'normalize_dict', 'width_determine',
           'single_sankey', 'sankeydiag', 'add_bars_to_plot', 'delta_text_plotter', 'delta_dots_plotter',
           'slopegraph_plotter', 'plot_minimeta_or_deltadelta_violins', 'effect_size_curve_plotter', 'gridkey_plotter',
           'barplotter', 'table_for_horizontal_plots', 'add_counts_to_prop_plots', 'swarmplot', 'SwarmPlot']

# %% ../nbs/API/plot_tools.ipynb 4
import math
import warnings
import itertools
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.axes as axes
import matplotlib.patches as mpatches
from collections import defaultdict
from typing import List, Tuple, Dict, Iterable, Union
from pandas.api.types import CategoricalDtype
from matplotlib.colors import ListedColormap

# %% ../nbs/API/plot_tools.ipynb 5
def halfviolin(v, half="right", fill_color="k", alpha=1, line_color="k", line_width=0):
    for b in v["bodies"]:
        V = b.get_paths()[0].vertices

        mean_vertical = np.mean(V[:, 0])
        mean_horizontal = np.mean(V[:, 1])

        if half == "right":
            V[:, 0] = np.clip(V[:, 0], mean_vertical, np.inf)
        elif half == "left":
            V[:, 0] = np.clip(V[:, 0], -np.inf, mean_vertical)
        elif half == "bottom":
            V[:, 1] = np.clip(V[:, 1], -np.inf, mean_horizontal)
        elif half == "top":
            V[:, 1] = np.clip(V[:, 1], mean_horizontal, np.inf)

        b.set_color(fill_color)
        b.set_alpha(alpha)
        b.set_edgecolor(line_color)
        b.set_linewidth(line_width)


def get_swarm_spans(coll):
    """
    Given a matplotlib Collection, will obtain the x and y spans
    for the collection. Will return None if this fails.
    """
    if coll is None:
        raise ValueError("The collection `coll` parameter cannot be None")

    x, y = np.array(coll.get_offsets()).T
    try:
        return x.min(), x.max(), y.min(), y.max()
    except ValueError as e:
        warnings.warn(f"Failed to calculate spans for the collection. Details: {e}")
        return None


def error_bar(
    data: pd.DataFrame,  # This DataFrame should be in 'long' format.
    x: str,  # x column to be plotted.
    y: str,  # y column to be plotted.
    type: str = "mean_sd",  # Choose from ['mean_sd', 'median_quartiles']. Plots the summary statistics for each group. If 'mean_sd', then the mean and standard deviation of each group is plotted as a gapped line. If 'median_quantiles', then the median and 25th and 75th percentiles of each group is plotted instead.
    offset: float = 0.2,  # Give a single float (that will be used as the x-offset of all gapped lines), or an iterable containing the list of x-offsets.
    ax=None,  # If a matplotlib Axes object is specified, the gapped lines will be plotted in order on this axes. If None, the current axes (plt.gca()) is used.
    line_color="black",  # The color of the gapped lines.
    gap_width_percent=1,  # The width of the gap in the gapped lines, as a percentage of the y-axis span.
    pos: list = [
        0,
        1,
    ],  # The positions of the error bars for the sankey_error_bar method.
    method: str = "gapped_lines",  # The method to use for drawing the error bars. Options are: 'gapped_lines', 'proportional_error_bar', and 'sankey_error_bar'.
    horizontal: bool = False,  # If True, the error bars will be horizontal. If False, the error bars will be vertical.
    **kwargs: dict,
):
    """
    Function to plot the standard deviations as vertical errorbars.
    The mean is a gap defined by negative space.

    This function combines the functionality of gapped_lines(),
    proportional_error_bar(), and sankey_error_bar().

    """

    if gap_width_percent < 0 or gap_width_percent > 100:
        raise ValueError("`gap_width_percent` must be between 0 and 100.")
    if method not in ["gapped_lines", "proportional_error_bar", "sankey_error_bar"]:
        raise ValueError(
            "Invalid `method`. Must be one of 'gapped_lines', \
                         'proportional_error_bar', or 'sankey_error_bar'."
        )

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

    if horizontal:
        ax_ylims = ax.get_xlim()
    else:
        ax_ylims = ax.get_ylim()
    ax_yspan = np.abs(ax_ylims[1] - ax_ylims[0])
    gap_width = ax_yspan * gap_width_percent / 100

    keys = kwargs.keys()
    if "clip_on" not in keys:
        kwargs["clip_on"] = False

    if "zorder" not in keys:
        kwargs["zorder"] = 5

    if "lw" not in keys:
        kwargs["lw"] = 2.0

    if isinstance(data[x].dtype, pd.CategoricalDtype):
        group_order = pd.unique(data[x]).categories
    else:
        group_order = pd.unique(data[x])

    means = data.groupby(x, observed=False)[y].mean().reindex(index=group_order)

    if method in ["proportional_error_bar", "sankey_error_bar"]:
        g = lambda x: np.sqrt(
            (np.sum(x) * (len(x) - np.sum(x))) / (len(x) * len(x) * len(x))
        )
        sd = data.groupby(x, observed=False)[y].apply(g)
    else:
        sd = data.groupby(x, observed=False)[y].std().reindex(index=group_order)

    lower_sd = means - sd
    upper_sd = means + sd

    if (lower_sd < ax_ylims[0]).any() or (upper_sd > ax_ylims[1]).any():
        kwargs["clip_on"] = True

    medians = data.groupby(x, observed=False)[y].median().reindex(index=group_order)
    quantiles = (
        data.groupby(x, observed=False)[y].quantile([0.25, 0.75]).unstack().reindex(index=group_order)
    )
    lower_quartiles = quantiles[0.25]
    upper_quartiles = quantiles[0.75]

    if type == "mean_sd":
        central_measures = means
        lows = lower_sd
        highs = upper_sd
    elif type == "median_quartiles":
        central_measures = medians
        lows = lower_quartiles
        highs = upper_quartiles
    else:
        raise ValueError("Only accepted values for type are ['mean_sd', 'median_quartiles']")

    n_groups = len(central_measures)

    if isinstance(line_color, str):
        custom_palette = np.repeat(line_color, n_groups)
    else:
        if len(line_color) != n_groups:
            err1 = "{} groups are being plotted, but ".format(n_groups)
            err2 = "{} colors(s) were supplied in `line_color`.".format(len(line_color))
            raise ValueError(err1 + err2)
        custom_palette = line_color

    try:
        len_offset = len(offset)
    except TypeError:
        offset = np.repeat(offset, n_groups)
        len_offset = len(offset)

    if len_offset != n_groups:
        err1 = "{} groups are being plotted, but ".format(n_groups)
        err2 = "{} offset(s) were supplied in `offset`.".format(len_offset)
        raise ValueError(err1 + err2)

    kwargs["zorder"] = kwargs["zorder"]

    for xpos, val in enumerate(central_measures.index):
        central_measure = central_measures[val]
        kwargs["color"] = custom_palette[xpos]

        if method == "sankey_error_bar":
            _xpos = pos[xpos] + offset[xpos]
        else:
            _xpos = xpos + offset[xpos]

        # Fix for the non-string x-axis issue #108
        if central_measures.index.dtype.name == "category":
            low = lows[xpos]
            high = highs[xpos]
        else: 
            low = lows[val]
            high = highs[val]

        if low == high == central_measure:
            if horizontal:
                low2mean_x, low2mean_y = [low, central_measure], [_xpos, _xpos]
                mean2high_x, mean2high_y = [central_measure, high], [_xpos, _xpos]
            else:
                low2mean_x, low2mean_y = [_xpos, _xpos], [low, central_measure]
                mean2high_x, mean2high_y = [_xpos, _xpos], [central_measure, high]
        else:
            if horizontal:
                low2mean_x, low2mean_y = [low, central_measure - gap_width], [_xpos, _xpos]
                mean2high_x, mean2high_y = [central_measure + gap_width, high], [_xpos, _xpos]
            else:
                low2mean_x, low2mean_y = [_xpos, _xpos], [low, central_measure - gap_width]
                mean2high_x, mean2high_y = [_xpos, _xpos], [central_measure + gap_width, high]
        # Add lines
        ax.add_line(mlines.Line2D(
                    low2mean_x, low2mean_y, **kwargs
                ))
        ax.add_line(mlines.Line2D(
                    mean2high_x, mean2high_y, **kwargs
                ))
        
def check_data_matches_labels(
    labels,  # list of input labels
    data,  # Pandas Series of input data
    side: str,  # 'left' or 'right' on the sankey diagram
):
    """
    Function to check that the labels and data match in the sankey diagram.
    And enforce labels and data to be lists.
    Raises an exception if the labels and data do not match.
    """
    if len(labels) > 0:
        if isinstance(data, list):
            data = set(data)
        if isinstance(data, pd.Series):
            data = set(data.unique())
        if isinstance(labels, list):
            labels = set(labels)
        if labels != data:
            msg = "\n"
            if len(labels) <= 20:
                msg = "Labels: " + ",".join(labels) + "\n"
            if len(data) < 20:
                msg += "Data: " + ",".join(data)
            raise Exception(f"{side} labels and data do not match.{msg}")


def normalize_dict(nested_dict, target):
    """
    Normalizes the values in a nested dictionary based on a target dictionary.

    This function iterates through a nested dictionary, calculates the sum of values for each key
    across all sub-dictionaries, and then normalizes these values according to a target dictionary.
    The normalization is performed such that the values in each sub-dictionary are proportionally
    scaled to match the corresponding 'right' values in the target dictionary.

    Parameters:
    nested_dict (dict of dict): A nested dictionary where each key maps to another dictionary.
                                The values in these inner dictionaries are subject to normalization.
    target (dict): A dictionary with the target values for normalization. Each key in nested_dict
                   should have a corresponding key in target, and each target[key] should be a
                   dictionary with a 'right' key containing the target normalization value.

    Returns:
    dict: The normalized nested dictionary. The original nested_dict is modified in place.

    Note:
    - If the sum of values for a particular key in nested_dict is zero, the normalized value is set to 0.
    - If a key in a sub-dictionary of nested_dict does not exist in the target dictionary, the
      corresponding 'right' value from the target dictionary is directly assigned.
    - The function modifies the input nested_dict in place and also returns it.
    """
    val = {}
    for key in nested_dict.keys():
        val[key] = np.sum(
            [
                nested_dict[sub_key][key]
                for sub_key in nested_dict.keys()
                if key in nested_dict[sub_key]
            ]
        )

    for key, value in nested_dict.items():
        if isinstance(value, dict):
            for subkey in value.keys():
                if subkey in val.keys():
                    if val[subkey] != 0:
                        # Address the problem when one of the labels has zero value
                        value[subkey] = (
                            value[subkey] * target[subkey]["right"] / val[subkey]
                        )
                    else:
                        value[subkey] = 0
                else:
                    value[subkey] = target[subkey]["right"]
    return nested_dict


def width_determine(labels, data, pos="left"):
    """
    Calculates normalized width positions for a set of labels based on their associated data.

    This function is designed to determine width positions for plotting or graphical representation.
    It takes into account the cumulative weight of each label in the data and adjusts their positions
    accordingly. The function allows for adjusting the position of labels to either the 'left' or 'right'.

    Parameters:
    labels (list): A list of labels whose width positions are to be calculated.
    data (DataFrame): A pandas DataFrame containing the data used for calculating width positions.
                      The DataFrame should have columns corresponding to the 'pos' and 'posWeight'.
    pos (str, optional): The position of labels. It can be either 'left' or 'right'. Defaults to 'left'.

    Returns:
    defaultdict: A dictionary where each key is a label and the value is another dictionary with keys
                 'bottom', 'top', and 'pos', representing the calculated width positions.

    Note:
    The function assumes that the data DataFrame contains columns named after the value of 'pos' and
    an additional column named 'posWeight' which represents the weight of each label.
    """
    if labels is None:
        raise ValueError("The `labels` parameter cannot be None")

    if data is None:
        raise ValueError("The `data` parameter cannot be None")
    
    widths_norm = defaultdict()
    for i, label in enumerate(labels):
        myD = {}
        myD[pos] = data[data[pos] == label][pos + "Weight"].sum()
        if len(labels) != 1:
            if i == 0:
                myD["bottom"] = 0
                myD[pos] -= 0.01
                myD["top"] = myD[pos]
            elif i == len(labels) - 1:
                myD[pos] -= 0.01
                myD["bottom"] = 1 - myD[pos]
                myD["top"] = 1
            else:
                myD[pos] -= 0.02
                myD["bottom"] = widths_norm[labels[i - 1]]["top"] + 0.02
                myD["top"] = myD["bottom"] + myD[pos]
        else:
            myD["bottom"] = 0
            myD["top"] = 1
        widths_norm[label] = myD
    return widths_norm


def single_sankey(
    left: np.array,  # data on the left of the diagram
    right: np.array,  # data on the right of the diagram, len(left) == len(right)
    xpos: float = 0,  # the starting point on the x-axis
    left_weight: np.array = None,  # weights for the left labels, if None, all weights are 1
    right_weight: np.array = None,  # weights for the right labels, if None, all weights are corresponding left_weight
    colorDict: dict = None,  # input format: {'label': 'color'}
    left_labels: list = None,  # labels for the left side of the diagram. The diagram will be sorted by these labels.
    right_labels: list = None,  # labels for the right side of the diagram. The diagram will be sorted by these labels.
    ax=None,  # matplotlib axes to be drawn on
    flow: bool = True,  # if True, draw the sankey in a flow, else draw 1 vs 1 Sankey diagram for each group comparison
    sankey: bool = True,  # if True, draw the sankey diagram, else draw barplot
    width=0.5,
    alpha=0.65,
    bar_width=0.2,
    error_bar_on: bool = True,  # if True, draw error bar for each group comparison
    strip_on: bool = True,  # if True, draw strip for each group comparison
    one_sankey: bool = False,  # if True, only draw one sankey diagram
    right_color: bool = False,  # if True, each strip of the diagram will be colored according to the corresponding left labels
    align: str = "center",  # if 'center', the diagram will be centered on each xtick,  if 'edge', the diagram will be aligned with the left edge of each xtick
    horizontal: bool = False,  # if True, the horizontal format for the sankey diagram will be used
):
    """
    Make a single Sankey diagram showing proportion flow from left to right

    Original code from: https://github.com/anazalea/pySankey
    
    Changes are added to normalize each diagram's height to be 1

    """

    # Initiating values
    if ax is None:
        ax = plt.gca()

    if left_weight is None:
        left_weight = []
    if right_weight is None:
        right_weight = []
    if left_labels is None:
        left_labels = []
    if right_labels is None:
        right_labels = []
    # Check weights
    if len(left_weight) == 0:
        left_weight = np.ones(len(left))
    if len(right_weight) == 0:
        right_weight = np.ones(len(right))

    # Create Dataframe
    if isinstance(left, pd.Series):
        left.reset_index(drop=True, inplace=True)
    if isinstance(right, pd.Series):
        right.reset_index(drop=True, inplace=True)
    dataFrame = pd.DataFrame(
        {
            "left": left,
            "right": right,
            "left_weight": left_weight,
            "right_weight": right_weight,
        },
        index=range(len(left)),
    )

    if dataFrame[["left", "right"]].isnull().any(axis=None):
        raise Exception("Sankey graph does not support null values.")

    # Identify all labels that appear 'left' or 'right'
    allLabels = pd.Series(
        np.sort(np.r_[dataFrame.left.unique(), dataFrame.right.unique()])[::-1]
    ).unique()

    # Identify left labels
    if len(left_labels) == 0:
        left_labels = pd.Series(np.sort(dataFrame.left.unique())[::-1]).unique()
    else:
        check_data_matches_labels(left_labels, dataFrame["left"], "left")

    # Identify right labels
    if len(right_labels) == 0:
        right_labels = pd.Series(np.sort(dataFrame.right.unique())[::-1]).unique()
    else:
        check_data_matches_labels(left_labels, dataFrame["right"], "right")

    # If no colorDict given, make one
    if colorDict is None:
        colorDict = {}
        palette = "hls"
        colorPalette = sns.color_palette(palette, len(allLabels))
        for i, label in enumerate(allLabels):
            colorDict[label] = colorPalette[i]
        fail_color = {0: "grey"}
        colorDict.update(fail_color)
    else:
        missing = [label for label in allLabels if label not in colorDict.keys()]
        if missing:
            msg = "The palette parameter is missing values for the following labels : "
            msg += "{}".format(", ".join(missing))
            raise ValueError(msg)

    if align not in ("center", "edge"):
        err = "{} assigned for `align` is not valid.".format(align)
        raise ValueError(err)
    
    if align == "center":
        try:
            leftpos = xpos - width / 2
        except TypeError as e:
            raise TypeError(
                f"the dtypes of parameters x ({xpos.dtype}) "
                f"and width ({width.dtype}) "
                f"are incompatible"
            ) from e
    else:
        leftpos = xpos

    # Combine left and right arrays to have a pandas.DataFrame in the 'long' format
    left_series = pd.Series(left, name="values").to_frame().assign(groups="left")
    right_series = pd.Series(right, name="values").to_frame().assign(groups="right")
    concatenated_df = pd.concat([left_series, right_series], ignore_index=True)

    # Determine positions of left label patches and total widths
    # We also want the height of the graph to be 1
    leftWidths_norm = defaultdict()
    for i, left_label in enumerate(left_labels):
        myD = {}
        myD["left"] = (
            dataFrame[dataFrame.left == left_label].left_weight.sum()
            / dataFrame.left_weight.sum()
        )
        if len(left_labels) != 1:
            if i == 0:
                myD["bottom"] = 0
                myD["left"] -= 0.01
                myD["top"] = myD["left"]
            elif i == len(left_labels) - 1:
                myD["left"] -= 0.01
                myD["bottom"] = 1 - myD["left"]
                myD["top"] = 1
            else:
                myD["left"] -= 0.02
                myD["bottom"] = leftWidths_norm[left_labels[i - 1]]["top"] + 0.02
                myD["top"] = myD["bottom"] + myD["left"]
                topEdge = myD["top"]
        else:
            myD["bottom"] = 0
            myD["top"] = 1
            myD["left"] = 1
        leftWidths_norm[left_label] = myD

    # Determine positions of right label patches and total widths
    rightWidths_norm = defaultdict()
    for i, right_label in enumerate(right_labels):
        myD = {}
        myD["right"] = (
            dataFrame[dataFrame.right == right_label].right_weight.sum()
            / dataFrame.right_weight.sum()
        )
        if len(right_labels) != 1:
            if i == 0:
                myD["bottom"] = 0
                myD["right"] -= 0.01
                myD["top"] = myD["right"]
            elif i == len(right_labels) - 1:
                myD["right"] -= 0.01
                myD["bottom"] = 1 - myD["right"]
                myD["top"] = 1
            else:
                myD["right"] -= 0.02
                myD["bottom"] = rightWidths_norm[right_labels[i - 1]]["top"] + 0.02
                myD["top"] = myD["bottom"] + myD["right"]
                topEdge = myD["top"]
        else:
            myD["bottom"] = 0
            myD["top"] = 1
            myD["right"] = 1
        rightWidths_norm[right_label] = myD

    # Total width of the graph
    xMax = width

    # Plot vertical bars for each label
    for left_label in left_labels:
        if horizontal:
            fill_method = ax.fill_betweenx
        else:
            fill_method = ax.fill_between
        fill_method(
                [leftpos + (-(bar_width) * xMax * 0.5), leftpos + (bar_width * xMax * 0.5)],
                2 * [leftWidths_norm[left_label]["bottom"]],
                2 * [leftWidths_norm[left_label]["top"]],
                color=colorDict[left_label],
                alpha=0.99,
            )
    if (not flow and sankey) or one_sankey:
        for right_label in right_labels:
            if horizontal:
                fill_method = ax.fill_betweenx
            else:
                fill_method = ax.fill_between
            fill_method(
                [
                    xMax + leftpos + (-bar_width * xMax * 0.5),
                    leftpos + xMax + (bar_width * xMax * 0.5),
                ],
                2 * [rightWidths_norm[right_label]["bottom"]],
                2 * [rightWidths_norm[right_label]["top"]],
                color=colorDict[right_label],
                alpha=0.99,
            )

    # Plot error bars
    if error_bar_on and strip_on:
        if horizontal:
            error_bar(
                concatenated_df,
                x="groups",
                y="values",
                ax=ax,
                offset=0,
                gap_width_percent=2,
                method="sankey_error_bar",
                pos=[leftpos, leftpos + xMax],
                horizontal=True,
            )
        else:
            error_bar(
                concatenated_df,
                x="groups",
                y="values",
                ax=ax,
                offset=0,
                gap_width_percent=2,
                method="sankey_error_bar",
                pos=[leftpos, leftpos + xMax],
            )

    # Determine widths of individual strips, all widths are normalized to 1
    ns_l = defaultdict()
    ns_r = defaultdict()
    ns_l_norm = defaultdict()
    ns_r_norm = defaultdict()
    for left_label in left_labels:
        leftDict = {}
        rightDict = {}
        for right_label in right_labels:
            leftDict[right_label] = dataFrame[
                (dataFrame.left == left_label) & (dataFrame.right == right_label)
            ].left_weight.sum()

            rightDict[right_label] = dataFrame[
                (dataFrame.left == left_label) & (dataFrame.right == right_label)
            ].right_weight.sum()
        factorleft = leftWidths_norm[left_label]["left"] / sum(leftDict.values())
        leftDict_norm = {k: v * factorleft for k, v in leftDict.items()}
        ns_l_norm[left_label] = leftDict_norm
        ns_r[left_label] = rightDict

    # ns_r should be using a different way of normalization to fit the right side
    # It is normalized using the value with the same key in each sub-dictionary
    ns_r_norm = normalize_dict(ns_r, rightWidths_norm)

    # Plot strips
    if sankey and strip_on:
        for left_label, right_label in itertools.product(left_labels, right_labels):
            labelColor = left_label
            
            if right_color:
                labelColor = right_label
            
            if len(dataFrame[(dataFrame.left == left_label) & 
                        (dataFrame.right == right_label)]) > 0:
                # Create array of y values for each strip, half at left value,
                # half at right, convolve
                ys_d = np.array(
                    50 * [leftWidths_norm[left_label]["bottom"]]
                    + 50 * [rightWidths_norm[right_label]["bottom"]]
                )
                ys_d = np.convolve(ys_d, 0.05 * np.ones(20), mode="valid")
                ys_d = np.convolve(ys_d, 0.05 * np.ones(20), mode="valid")
                # to remove the array wrapping behaviour of black
                # fmt: off
                ys_u = np.array(50 * [leftWidths_norm[left_label]['bottom'] + ns_l_norm[left_label][right_label]] + \
                    50 * [rightWidths_norm[right_label]['bottom'] + ns_r_norm[left_label][right_label]])
                # fmt: on
                ys_u = np.convolve(ys_u, 0.05 * np.ones(20), mode="valid")
                ys_u = np.convolve(ys_u, 0.05 * np.ones(20), mode="valid")

                # Update bottom edges at each label so next strip starts at the right place
                leftWidths_norm[left_label]["bottom"] += ns_l_norm[left_label][right_label]
                rightWidths_norm[right_label]["bottom"] += ns_r_norm[left_label][
                    right_label
                ]
                if horizontal:
                    fill_method = ax.fill_betweenx
                else:
                    fill_method = ax.fill_between
                fill_method(
                    np.linspace(
                        leftpos + (bar_width * xMax * 0.5),
                        leftpos + xMax - (bar_width * xMax * 0.5),
                        len(ys_d),
                    ),
                    ys_d,
                    ys_u,
                    alpha=alpha,
                    color=colorDict[labelColor],
                    edgecolor="none",
                )

def sankeydiag(
    data: pd.DataFrame,
    xvar: str,  # x column to be plotted.
    yvar: str,  # y column to be plotted.
    temp_all_plot_groups: list,
    idx: list,
    temp_idx: list,
    left_labels: list = None,  # labels for the left side of the diagram. The diagram will be sorted by these labels.
    right_labels: list = None,  # labels for the right side of the diagram. The diagram will be sorted by these labels.
    palette: str | dict = None,
    ax=None,  # matplotlib axes to be drawn on
    flow: bool = True,  # if True, draw the sankey in a flow, else draw 1 vs 1 Sankey diagram for each group comparison
    sankey: bool = True,  # if True, draw the sankey diagram, else draw barplot
    one_sankey: bool = False,  # determined by the driver function on plotter.py, if True, draw the sankey diagram across the whole raw data axes
    width: float = 0.4,  # the width of each sankey diagram
    right_color: bool = False,  # if True, each strip of the diagram will be colored according to the corresponding left labels
    align: str = "center",  # the alignment of each sankey diagram, can be 'center' or 'left'
    alpha: float = 0.65,  # the transparency of each strip
    horizontal: bool = False,  # if True, the horizontal format for the sankey diagram will be used
    **kwargs,
):
    """
    Read in melted pd.DataFrame, and draw multiple sankey diagram on a single axes
    using the value in column yvar according to the value in column xvar
    left_idx in the column xvar is on the left side of each sankey diagram
    right_idx in the column xvar is on the right side of each sankey diagram

    """
    if "width" in kwargs:
        width = kwargs["width"]

    if "align" in kwargs:
        align = kwargs["align"]

    if "alpha" in kwargs:
        alpha = kwargs["alpha"]

    if "right_color" in kwargs:
        right_color = kwargs["right_color"]

    if "bar_width" in kwargs:
        bar_width = kwargs["bar_width"]

    if "sankey" in kwargs:
        sankey = kwargs["sankey"]

    if "flow" in kwargs:
        flow = kwargs["flow"]

    fontsize = kwargs.pop("fontsize")

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

    left_idx = []
    right_idx = []
    # Design for Sankey Flow Diagram
    sankey_idx = (
    [
        (control, test)
        for i in idx
        for control, test in zip(
            i[:],
            (tuple(i[1:]) + (i[0],)) if isinstance(i, tuple) else (list(i[1:]) + [i[0]])
        )
    ]
    if flow
    else temp_idx   
    )

    for i in sankey_idx:
        left_idx.append(i[0])
        right_idx.append(i[1])

    if len(temp_all_plot_groups) == 2:
        one_sankey = True
        left_idx.pop()
        right_idx.pop()  # Remove the last element from two lists

    # two_col_sankey = True if proportional == True and one_sankey == False and sankey == True and flow == False else False

    allLabels = pd.Series(np.sort(data[yvar].unique())[::-1]).unique()

    # Check if all the elements in left_idx and right_idx are in xvar column
    unique_xvar = data[xvar].unique()
    if not all(elem in unique_xvar for elem in left_idx):
        raise ValueError(f"{left_idx} not found in {xvar} column")
    if not all(elem in unique_xvar for elem in right_idx):
        raise ValueError(f"{right_idx} not found in {xvar} column")

    xpos = 0

    # For baseline comparison, broadcast left_idx to the same length as right_idx
    # so that the left of sankey diagram will be the same
    # For sequential comparison, left_idx and right_idx can have anything different
    # but should have the same length
    if len(left_idx) == 1:
        broadcasted_left = np.broadcast_to(left_idx, len(right_idx))
    elif len(left_idx) != len(right_idx):
        raise ValueError(f"left_idx and right_idx should have the same length")
    else:
        broadcasted_left = left_idx

    if isinstance(palette, dict):
        if not all(key in allLabels for key in palette.keys()):
            raise ValueError(f"keys in palette should be in {yvar} column")
        plot_palette = palette
    elif isinstance(palette, str):
        plot_palette = {}
        colorPalette = sns.color_palette(palette, len(allLabels))
        for i, label in enumerate(allLabels):
            plot_palette[label] = colorPalette[i]
    else:
        plot_palette = None

    # Create a strip_on list to determine whether to draw the strip during repeated measures
    strip_on = [
        int(right not in broadcasted_left[:i]) for i, right in enumerate(right_idx)
    ]

    draw_idx = list(zip(broadcasted_left, right_idx))
    for i, (left, right) in enumerate(draw_idx):
        if not one_sankey:
            if flow:
                width = 1
                align = "edge"
                sankey = (
                    False if i == len(draw_idx) - 1 else sankey
                )  # Remove last strip in flow
            error_bar_on = (
                False if i == len(draw_idx) - 1 and flow else True
            )  # Remove last error_bar in flow
            bar_width = 0.4 if sankey == False and flow == False else bar_width
            single_sankey(
                data[data[xvar] == left][yvar],
                data[data[xvar] == right][yvar],
                xpos=xpos,
                ax=ax,
                colorDict=plot_palette,
                width=width,
                left_labels=left_labels,
                right_labels=right_labels,
                strip_on=strip_on[i],
                right_color=right_color,
                bar_width=bar_width,
                sankey=sankey,
                error_bar_on=error_bar_on,
                flow=flow,
                align=align,
                alpha=alpha,
                horizontal=horizontal,
            )
            xpos += 1
        else:
            xpos = 0
            width = 1
            if not sankey:
                bar_width = 0.5
            single_sankey(
                data[data[xvar] == left][yvar],
                data[data[xvar] == right][yvar],
                xpos=xpos,
                ax=ax,
                colorDict=plot_palette,
                width=width,
                left_labels=left_labels,
                right_labels=right_labels,
                right_color=right_color,
                bar_width=bar_width,
                sankey=sankey,
                one_sankey=one_sankey,
                flow=False,
                align="edge",
                alpha=alpha,
                horizontal=horizontal,
            )

    # Now only draw vs xticks for two-column sankey diagram

    if not one_sankey or (sankey and not flow):
        sankey_tick_vals = (
            [f"{left}" for left in broadcasted_left]
            if flow
            else [f"{left} v.s. {right}" if horizontal
                else f"{left}\n v.s.\n{right}"
                for left, right in zip(broadcasted_left, right_idx)
            ]
        )
        sankey_tick_locs = np.arange(len(right_idx))
    else:
        sankey_tick_vals, sankey_tick_locs = [broadcasted_left[0], right_idx[0]], [0, 1]

    if horizontal:
        ax.set_yticks(sankey_tick_locs)
        ax.set_yticklabels(sankey_tick_vals, fontsize = fontsize)
    else:
        ax.set_xticks(sankey_tick_locs)
        ax.set_xticklabels(sankey_tick_vals, fontsize = fontsize)

    return (left_idx, right_idx)

def add_bars_to_plot(bar_dict: dict, ax: axes.Axes, bar_kwargs: dict):
    """
    Add bars to the relevant axes.

    Parameters
    ----------
    bar_dict : dict
        Dictionary of bar values.
    ax : axes.Axes
        Matplotlib axis object to plot on.
    bar_kwargs : dict
        Keyword arguments for the bars.
    """
    og_xlim, og_ylim = ax.get_xlim(), ax.get_ylim()

    x_start_values, x_distances, y_start_values, y_distances, colors = bar_dict.values()

    for start_x, start_y, distance_x, distance_y, current_color in zip(
                                                                        x_start_values, 
                                                                        y_start_values, 
                                                                        x_distances, 
                                                                        y_distances, 
                                                                        colors
                                                                    ):
        ax.add_patch(mpatches.Rectangle((start_x, start_y), 
                                        distance_x, distance_y, 
                                        color=current_color, **bar_kwargs
                                        )
                    )
    ax.set_xlim(og_xlim)
    ax.set_ylim(og_ylim) 

def delta_text_plotter(
        results: pd.DataFrame, 
        ax_to_plot: object, 
        ticks_to_plot: list, 
        delta_text_kwargs: dict, 
        color_col: str, 
        plot_palette_raw: dict, 
        show_pairs: bool,
        float_contrast: bool,
        extra_delta: float,
        bootstraps_color_by_group: bool = False
    ):
    """
    Add delta text to the contrast plot.

    Parameters
    ----------
    results : object (Dataframe)
        Dataframe of contrast object comparisons.
    ax_to_plot : axes.Axes
        Matplotlib axis object to plot on.
    ticks_to_plot : list
        List of indices of the contrast objects.
    delta_text_kwargs : dict
        Keyword arguments for the delta text.
    color_col : str
        Column name of the color column.
    plot_palette_raw : dict
        Dictionary of colors used in the plot.
    show_pairs : bool
        Whether the data is paired and show pairs.
    float_contrast : bool
        Whether the DABEST plot uses Gardner-Altman or Cummings.
    extra_delta : float or None
        The extra mini-meta or delta-delta value if applicable.
    bootstraps_color_by_group : bool, optional
        Whether to color the bootstraps by group. Default is False.
    """
    # Colors
    from .misc_tools import color_picker
    delta_text_colors = color_picker(color_type = 'delta_text',
                                     kwargs = delta_text_kwargs, 
                                     elements = ticks_to_plot, 
                                     color_col = color_col, 
                                     show_pairs = show_pairs, 
                                     color_palette = plot_palette_raw,
                                     bootstraps_color_by_group = bootstraps_color_by_group
                                    )

    num_of_elements = len(ticks_to_plot) + 1 if extra_delta is not None else len(ticks_to_plot)

    # Collect the means for the delta text
    delta_values = []
    for j, tick in enumerate(ticks_to_plot):
        delta_values.append(results.difference[int(j)])
    if extra_delta is not None: 
        delta_values.append(extra_delta)
        delta_text_colors.append('black')

    # Collect the X-coordinates for the delta text
    delta_text_x_coordinates = delta_text_kwargs.pop('x_coordinates')
    delta_text_x_offset = delta_text_kwargs.pop('offset')

    if delta_text_x_coordinates is not None:
        if not isinstance(delta_text_x_coordinates, (list, tuple)) or not all(isinstance(x, (int, float)) for x in delta_text_x_coordinates):
            raise TypeError("delta_text_kwargs['x_coordinates'] must be a list of x-coordinates.")
        if len(delta_text_x_coordinates) != num_of_elements:
            raise ValueError("delta_text_kwargs['x_coordinates'] must have the same length as the number of ticks to plot.")
    else:
        x_adjust = (-0.4 if float_contrast else 0.48) + delta_text_x_offset
        delta_text_x_coordinates = [x+x_adjust for x in ticks_to_plot]
        if extra_delta is not None: delta_text_x_coordinates.append(max(ticks_to_plot)+1+x_adjust)

    # Collect the Y-coordinates for the delta text
    delta_text_y_coordinates = delta_text_kwargs.pop('y_coordinates')
    if float_contrast: delta_text_kwargs["va"] = 'bottom' if results.difference[0] >= 0 else 'top'

    if delta_text_y_coordinates is not None:
        if not isinstance(delta_text_y_coordinates, (list, tuple)) or not all(isinstance(y, (int, float)) for y in delta_text_y_coordinates):
            raise TypeError("delta_text_kwargs['y_coordinates'] must be a list of y-coordinates.")
        if len(delta_text_y_coordinates) != num_of_elements:
            raise ValueError("delta_text_kwargs['y_coordinates'] must have the same length as the number of ticks to plot.")
    else:
        delta_text_y_coordinates = delta_values

    # Plot the delta text
    for x, y, text, color in zip(delta_text_x_coordinates, delta_text_y_coordinates, delta_values, delta_text_colors):
        delta_text = np.format_float_positional(text, precision=2, sign=True, trim="k", min_digits=2)
        ax_to_plot.text(x, y, delta_text, color=color, zorder=5, **delta_text_kwargs)

def delta_dots_plotter(
        plot_data: pd.DataFrame, 
        contrast_axes: axes.Axes, 
        delta_id_col: str, 
        idx: list, 
        xvar: str, 
        yvar: str, 
        is_paired: bool, 
        color_col: str, 
        float_contrast: bool, 
        plot_palette_raw: dict, 
        delta_dot_kwargs: dict, 
        horizontal: bool
    ):
    """
    Parameters
    ----------
    plot_data : object (Dataframe)
        Dataframe of the plot data.
    contrast_axes : axes.Axes
        Matplotlib axis object to plot on.
    delta_id_col : str
        Column name of the delta id column.
    idx : list
        List of indices of the contrast objects.
    xvar : str
        Column name of the x variable.
    yvar : str
        Column name of the y variable.
    is_paired : bool
        Whether the data is paired.
    color_col : str
        Column name of the color column.
    float_contrast : bool
        Whether the DABEST plot uses Gardner-Altman or Cummings
    plot_palette_raw : dict
        Dictionary of colors used in the plot.
    delta_dot_kwargs : dict
        Keyword arguments for the delta dots.
    horizontal : bool
        If the rawplot is horizontal.
    """

    # Checks and initializations
    # from .plot_tools import swarmplot
    delta_dot_color = delta_dot_kwargs.pop('color')
    if color_col is not None:
        plot_palette_deltapts = plot_palette_raw
        delta_plot_data = plot_data[[xvar, yvar, delta_id_col, color_col]]
    else:
        plot_palette_deltapts = delta_dot_color
        delta_plot_data = plot_data[[xvar, yvar, delta_id_col]]

    # TODO: to make jitter value more accurate and not just a hardcoded eyeball value
    jitter = 0.6 if float_contrast else 1 

    # Create dataframe of delta values
    final_deltas = pd.DataFrame()
    for i in idx:
        for j in i:
            if i.index(j) != 0:
                temp_df_exp = delta_plot_data[
                    delta_plot_data[xvar].str.contains(j)
                ].reset_index(drop=True)
                if is_paired == "baseline":
                    temp_df_cont = delta_plot_data[
                        delta_plot_data[xvar].str.contains(i[0])
                    ].reset_index(drop=True)
                elif is_paired == "sequential":
                    temp_df_cont = delta_plot_data[
                        delta_plot_data[xvar].str.contains(
                            i[i.index(j) - 1]
                        )
                    ].reset_index(drop=True)
                delta_df = temp_df_exp.copy()
                delta_df[yvar] = temp_df_exp[yvar] - temp_df_cont[yvar]
                final_deltas = pd.concat([final_deltas, delta_df])

    if horizontal:
        delta_dot_kwargs.update({'side': 'left'})
    # Plot the delta dots
    swarmplot(
        data=final_deltas,
        x=xvar,
        y=yvar,
        ax=contrast_axes,
        order=None,
        hue=color_col,
        palette=plot_palette_deltapts,
        jitter=jitter,
        is_drop_gutter=True,
        gutter_limit=1,
        horizontal=horizontal,
        **delta_dot_kwargs
    )
    contrast_axes.legend().set_visible(False)


def slopegraph_plotter(
        dabest_obj: object, 
        plot_data: pd.DataFrame, 
        xvar: str, 
        yvar: str, 
        color_col: str, 
        plot_palette_raw: dict, 
        slopegraph_kwargs: dict, 
        rawdata_axes: axes.Axes, 
        ytick_color: str, 
        temp_idx: list, 
        horizontal: bool,
        temp_all_plot_groups: list,
        plot_kwargs: dict,
        group_summaries_kwargs: dict
    ):
    """
    Add slopegraph to the rawdata axes.

    Parameters
    ----------
    dabest_obj : object
        DABEST object.
    plot_data : object (Dataframe)
        Dataframe of the plot data.
    xvar : str
        Column name of the x variable.
    yvar : str
        Column name of the y variable.
    color_col : str
        Column name of the color column.
    plot_palette_raw : dict
        Dictionary of colors used in the plot.
    slopegraph_kwargs : dict
        Keyword arguments for the slopegraph.
    rawdata_axes : axes.Axes
        Matplotlib axis object to plot on.
    ytick_color : str
        Color of the yticks.
    temp_idx : list
        List of indices of the contrast objects.
    horizontal : bool
        If the plotting will be in horizontal format.
    temp_all_plot_groups : list
        List of all plot groups.
    plot_kwargs : dict
        Keyword arguments for the plot.
    group_summaries_kwargs : dict, optional
        Keyword arguments for group summaries, if applicable.

    """
    # Jitter Kwargs 
    # With help from GitHub user: devMJBL
    jitter = slopegraph_kwargs.pop("jitter")
    if jitter > 1:
        err0 = "Jitter value is too high. Defaulting to 1."
        warnings.warn(err0)
        jitter = 1
    rng = np.random.default_rng(slopegraph_kwargs.pop("jitter_seed"))

    # Pivot the long (melted) data.
    if color_col is None:
        pivot_values = [yvar]
    else:
        pivot_values = [yvar, color_col]
    pivoted_plot_data = pd.pivot(
        data=plot_data,
        index=dabest_obj.id_col,
        columns=xvar,
        values=pivot_values,
    )

    x_start = 0
    for ii, current_tuple in enumerate(temp_idx):
        current_pair = pivoted_plot_data.loc[
            :, pd.MultiIndex.from_product([pivot_values, current_tuple])
        ].dropna()

        # Check for correct pairing
        if len(current_pair) == 0:
            raise ValueError('There are no pairs to plot... check original dataframe for correct ID pairing')

        current_pair = pivoted_plot_data.loc[
            :, pd.MultiIndex.from_product([pivot_values, current_tuple])
        ]
        grp_count = len(current_tuple)

        # Iterate through the data for the current tuple.
        for ID, observation in current_pair.iterrows():
            x_points = [t + 0.15*jitter*rng.standard_t(df=6, size=None) for t in range(x_start, x_start + grp_count)] # devMJBL
            y_points = observation[yvar].tolist()

            if color_col is None:
                slopegraph_kwargs["color"] = ytick_color
            else:
                color_key = observation[color_col].iloc[0]
                if isinstance(color_key, (str, np.int64, np.float64)):
                    slopegraph_kwargs["color"] = plot_palette_raw[color_key]
                    slopegraph_kwargs["label"] = color_key

            x_points, y_points = (y_points, x_points) if horizontal else (x_points, y_points)
            rawdata_axes.plot(x_points, y_points, **slopegraph_kwargs)

        # Add the group summaries if applicable.
        group_summaries = plot_kwargs.get("group_summaries", None)
        if group_summaries is not None:
            for key in ['gap_width_percent', 'offset']:
                group_summaries_kwargs.pop(key, None)
            group_summaries_kwargs['color'] = 'black' if group_summaries_kwargs.get('color') is None else group_summaries_kwargs['color']
            group_summaries_kwargs['capsize'] = 0 if group_summaries_kwargs.get('capsize') is None else group_summaries_kwargs['capsize']

            index_points = [t for t in range(x_start, x_start + grp_count)]
            av_points, err_points, lo_points, hi_points = [], [], [], []
            for group in range(len(index_points)):
                if group_summaries == "mean_sd":
                    av_points.append(current_pair.iloc[:, int(group)].mean())
                    err_points.append(current_pair.iloc[:, int(group)].std())
                elif group_summaries == "median_quartiles":
                    median = current_pair.iloc[:, int(group)].median()
                    av_points.append(median)
                    lo_points.append(median - current_pair.iloc[:, int(group)].quantile(0.25))
                    hi_points.append(current_pair.iloc[:, int(group)].quantile(0.75) - median)

            if group_summaries == "median_quartiles":
                err_points = [lo_points, hi_points] 

            # Plot the lines
            if horizontal:
                rawdata_axes.errorbar(
                        av_points,
                        index_points, 
                        xerr=err_points, 
                        **group_summaries_kwargs
                        )
            else:
                rawdata_axes.errorbar(
                        index_points, 
                        av_points, 
                        yerr=err_points, 
                        **group_summaries_kwargs
                        )

        x_start = x_start + grp_count

    # Set the tick labels, because the slopegraph plotting doesn't.
    if horizontal:
        rawdata_axes.set_yticks(np.arange(0, len(temp_all_plot_groups)))
        rawdata_axes.set_yticklabels(temp_all_plot_groups, fontsize = plot_kwargs.get("fontsize_rawxlabel"))
    else:
        rawdata_axes.set_xticks(np.arange(0, len(temp_all_plot_groups)))
        rawdata_axes.set_xticklabels(temp_all_plot_groups, fontsize = plot_kwargs.get("fontsize_rawxlabel"))
    

def plot_minimeta_or_deltadelta_violins(
        dabest_obj: object,
        type: str,
        ci_type: str, 
        rawdata_axes: axes.Axes,
        contrast_axes: axes.Axes, 
        contrast_kwargs: dict, 
        contrast_xtick_labels: list, 
        effect_size: str, 
        plot_kwargs: dict, 
        horizontal: bool, 
        show_pairs: bool,
        contrast_marker_kwargs: dict, 
        contrast_errorbar_kwargs: dict,
    ):
    """
    Add mini meta-analysis or delta-delta violin plots to the contrast plot.

    Parameters
    ----------
    dabest_obj : object
        DABEST Effectsize object delta-delta or mini_meta
    type: str
        mini_meta or delta_delta
    ci_type : str
        Type of confidence interval to plot.
    rawdata_axes : axes.Axes
        Matplotlib axis object to plot on.
    contrast_axes : axes.Axes
        Matplotlib axis object to plot on.
    contrast_kwargs : dict
        Keyword arguments for the violinplot.
    contrast_xtick_labels : list
        List of xtick labels for the contrast plot.
    effect_size : str
        Type of effect size to plot.
    plot_kwargs : dict
        Keyword arguments for the plot.
    horizontal : bool
        If the plot is horizontal.
    show_pairs : bool
        Whether the data is paired and shown in pairs.
    contrast_marker_kwargs: dict
        Keyword arguments for the effectsize marker.
    contrast_errorbar_kwargs: dict
        Keyword arguments for the effectsize errorbar.
    """

    # Plot the curve
    def extract_curve_data(dabest_object):
        try:
            data = dabest_object.bootstraps_weighted_delta
        except AttributeError:
            data = dabest_object.bootstraps_delta_delta

        ci_low, ci_high = dabest_object.results.get(ci_type+'_low')[0], dabest_object.results.get(ci_type+'_high')[0]
        return data, dabest_object.difference, ci_low, ci_high

    data, difference, ci_low, ci_high = extract_curve_data(dabest_obj)

    if contrast_kwargs.get('alpha') is not None:
        contrast_alpha = contrast_kwargs.pop('alpha')

    if horizontal:  
        contrast_kwargs.update({'orientation': 'horizontal', 'widths': 1})
        position = max(rawdata_axes.get_yticks()) + 1
        half = "bottom"
        effsize_x, effsize_y = difference, [position]
        ci_x, ci_y = [ci_low, ci_high], [position, position]
    else:
        position = max(rawdata_axes.get_xticks()) + 1
        half = "right"
        effsize_x, effsize_y = [position], difference
        ci_x, ci_y = [position, position], [ci_low, ci_high]

    v = contrast_axes.violinplot(
        data[~np.isinf(data)], positions=[position], **contrast_kwargs
        )

    halfviolin(v, fill_color="grey", alpha=contrast_alpha, half=half)

    # Plot the effect size.
    contrast_axes.plot(
        effsize_x,
        effsize_y,
        **contrast_marker_kwargs
    )
    # Plot the confidence interval.
    contrast_axes.plot(
        ci_x,
        ci_y,
        **contrast_errorbar_kwargs
    )

    # Add labels and ticks
    if horizontal:
        current_ylabels = rawdata_axes.get_yticklabels()
        if type == 'mini_meta':
            current_ylabels.extend(["Weighted delta"])
        elif effect_size == "hedges_g":
            current_ylabels.extend(["Delta g"])
        else:
            current_ylabels.extend(["Delta-delta"])

        rawdata_axes.set_yticks(np.append(rawdata_axes.get_yticks(), position))
        rawdata_axes.set_yticklabels(current_ylabels)
    else:
        if type == 'mini_meta':
            if show_pairs:
                contrast_xtick_labels.extend(["Weighted\n delta"])
            else:
                contrast_xtick_labels.extend(["Weighted delta"])
        elif effect_size == "hedges_g":
            contrast_xtick_labels.extend(["Delta g"])
        else:
            contrast_xtick_labels.extend(["Delta-delta"])

    # Create the delta-delta axes.
    if type == 'delta_delta' and not horizontal:
        if plot_kwargs["delta2_label"] is not None:
            delta2_label = plot_kwargs["delta2_label"]
        elif effect_size == "mean_diff":
            delta2_label = "Delta-delta"
        else:
            delta2_label = "Delta g"
        fontsize_delta2label = plot_kwargs["fontsize_delta2label"]
        delta2_axes = contrast_axes.twinx()
        delta2_axes.set_frame_on(False)
        delta2_axes.set_ylabel(delta2_label, fontsize=fontsize_delta2label)
        og_xlim_delta, og_ylim_delta = contrast_axes.get_xlim(), contrast_axes.get_ylim()
        delta2_axes.set_ylim(og_ylim_delta)
    else:
        delta2_axes = None

    return delta2_axes, contrast_xtick_labels


def effect_size_curve_plotter(
        ticks_to_plot: list, 
        ticks_for_baseline_ec: list, 
        results: pd.DataFrame, 
        ci_type: str, 
        contrast_axes: axes.Axes, 
        contrast_kwargs: dict, 
        bootstraps_color_by_group: bool, 
        plot_palette_contrast: dict,
        horizontal: bool, 
        contrast_marker_kwargs: dict, 
        contrast_errorbar_kwargs: dict,
        idx: list, 
        is_paired: bool, 
        contrast_paired_lines: bool, 
        contrast_paired_lines_kwargs: dict,
        show_baseline_ec: bool = False
    ):
    """
    Add effect size curves to the contrast plot.

    Parameters
    ----------
    ticks_to_plot : list
        List of indices of the contrast objects.
    ticks_for_baseline_ec : list
        List of indices of the baseline effect curve objects.
    results : object (Dataframe)
        Dataframe of contrast object comparisons.
    ci_type : str
        Type of confidence interval to plot.
    contrast_axes : axes.Axes
        Matplotlib axis object to plot on.
    contrast_kwargs : dict
        Keyword arguments for the violinplot.
    bootstraps_color_by_group : bool
        Whether to color the bootstraps by group.
    plot_palette_contrast : dict
        Dictionary of colors used in the contrast plot.
    horizontal : bool
        If the plot is horizontal.
    contrast_marker_kwargs: dict
        Keyword arguments for the effectsize marker.
    contrast_errorbar_kwargs: dict
        Keyword arguments for the effectsize errorbar.
    idx : list
        List of indices of the raw groups.
    is_paired : bool
        Whether the data is paired.
    contrast_paired_lines : bool
        Whether to add lines for repeated measures data.
    contrast_paired_lines_kwargs : dict
        Keyword arguments for the repeated measures lines.
    show_baseline_ec : bool
        Whether to show the baseline effect curve.
    """

    def plot_effect_size(tick, group, control, bootstrap, effsize, ci_low, ci_high):
        # Create the violinplot
        if horizontal:  
            contrast_kwargs.update({'orientation': 'horizontal', 'widths': 1})
            
        v = contrast_axes.violinplot(
            bootstrap[~np.isinf(bootstrap)],
            positions=[tick],
            **contrast_kwargs
        )
        
        # Color the violin plot
        fc = plot_palette_contrast[group] if bootstraps_color_by_group else "grey"
        half = "bottom" if horizontal else "right"
        halfviolin(v, fill_color=fc, alpha=contrast_alpha, half=half)

        # Plot the confidence interval
        if horizontal:
            ci_x, ci_y = [ci_low, ci_high], [tick, tick]
        else:
            ci_x, ci_y = [tick, tick], [ci_low, ci_high]
            
        contrast_axes.plot(ci_x, ci_y, **contrast_errorbar_kwargs)
        
        return "{}\nminus\n{}".format(group, control)
    
    if contrast_kwargs.get('alpha') is not None:
        contrast_alpha = contrast_kwargs.pop('alpha')

    # Plot the curves
    contrast_xtick_labels = []
    for j, tick in enumerate(ticks_to_plot):
        current_group = results.test[int(j)]
        current_control = results.control[int(j)]
        current_bootstrap = results.bootstraps[int(j)]
        current_effsize = results.difference[int(j)]
        current_ci_low = results.get(ci_type+'_low')[int(j)]
        current_ci_high = results.get(ci_type+'_high')[int(j)]

        # Plot the effect size marker
        if horizontal:
            effsize_x, effsize_y = current_effsize, [tick]
        else:
            effsize_x, effsize_y = [tick], current_effsize

        contrast_axes.plot(
            effsize_x,
            effsize_y,
            **contrast_marker_kwargs
        )

        label = plot_effect_size(tick, current_group, current_control, current_bootstrap,
                               current_effsize, current_ci_low, current_ci_high)
        contrast_xtick_labels.append(label)

    # Add baseline effect curve plotting
    bec_results = results.drop_duplicates(subset='control', keep='first').reset_index(drop=True)
    for j, tick in enumerate(ticks_for_baseline_ec):
        bec_group = bec_results.control[j]
        bec_control = bec_results.control[j]
        bec_bootstrap = bec_results.bec_bootstraps[j]
        bec_effsize = bec_results.bec_difference[j]
        bec_ci_low = bec_results.get('bec_'+ci_type+'_low')[j]
        bec_ci_high = bec_results.get('bec_'+ci_type+'_high')[j]
        
        # Plot the effect size marker regardless of show_baseline_ec
        if horizontal:
            effsize_x, effsize_y = bec_effsize, [tick]
        else:
            effsize_x, effsize_y = [tick], bec_effsize
            
        contrast_axes.plot(effsize_x, effsize_y, **contrast_marker_kwargs)
        
        if show_baseline_ec:
            _ = plot_effect_size(tick, bec_group, bec_control, bec_bootstrap, 
                               bec_effsize, bec_ci_low, bec_ci_high)
            # Baseline Curve doesn't need tick text

    # Add lines for repeated measures data
    if is_paired and contrast_paired_lines:
        temp_num = 0
        lines_to_plot_list = []

        for group in idx:
            new_group = []
            if len(group) >= 2:
                new_group.append(temp_num)
                for i in range(1, len(group)):
                    new_group.append(temp_num+i)
            temp_num += len(group)
            lines_to_plot_list.append(new_group)

        for group in lines_to_plot_list:
            if len(group) > 0:
                mean_diffs_for_lines = []
                for ticks in group:
                    if ticks in ticks_to_plot:
                        mean_diffs_for_lines.append(results.loc[ticks_to_plot.index(ticks)]["difference"])
                    else:
                        mean_diffs_for_lines.append(int(0))

                x_data = mean_diffs_for_lines if horizontal else group
                y_data = group if horizontal else mean_diffs_for_lines

                contrast_axes.plot(
                    x_data, 
                    y_data,
                    **contrast_paired_lines_kwargs
                )

    contrast_kwargs['alpha'] = contrast_alpha
    return current_group, current_control, current_effsize, contrast_xtick_labels

def gridkey_plotter(
        is_paired: bool, 
        idx: list,
        all_plot_groups: list, 
        gridkey: list, 
        rawdata_axes: axes.Axes, 
        contrast_axes: axes.Axes, 
        plot_data: pd.DataFrame, 
        xvar: str, 
        yvar: str, 
        results: pd.DataFrame, 
        show_delta2: bool, 
        show_mini_meta: bool, 
        x1_level: list, 
        experiment_label: list, 
        float_contrast: bool, 
        horizontal: bool, 
        delta_delta: object, 
        mini_meta: object, 
        effect_size: str, 
        gridkey_kwargs: dict,
    ):
    """
    Add gridkey to the contrast plot.

    Parameters
    ----------
    is_paired : bool
        Whether the data is paired.
    idx : list
        List of indices of the contrast objects.
    all_plot_groups : list
        List of all plot groups.
    gridkey : list
        List of gridkey rows.
    rawdata_axes : axes.Axes
        Matplotlib axis object for the raw data.
    contrast_axes : axes.Axes
        Matplotlib axis object for the contrast data.
    plot_data : object (Dataframe)
        Dataframe of the plot data.
    xvar : str
        Column name of the x variable.
    yvar : str
        Column name of the y variable.
    results : object (Dataframe)
        Dataframe of contrast object comparisons.
    show_delta2 : bool
        Whether to show the delta-delta.
    show_mini_meta : bool
        Whether to show the mini meta-analysis.
    x1_level : list
        List of x1 levels.
    experiment_label : list
        List of experiment labels.
    float_contrast : bool
        Whether the DABEST plot uses Gardner-Altman or Cummings
    horizontal : bool
        If the plot is horizontal.
    delta_delta : object
        delta-delta object.
    mini_meta : object
        Mini meta-analysis object.
    effect_size : str
        Type of effect size to plot
    gridkey_kwargs : dict
        Keyword arguments for the gridkey.
    """
    # Extract relevant kwargs
    gridkey_show_Ns = gridkey_kwargs["show_Ns"]
    gridkey_show_es = gridkey_kwargs["show_es"]
    gridkey_merge_pairs = gridkey_kwargs["merge_pairs"]
    gridkey_marker = gridkey_kwargs["marker"]
    gridkey_delimiters = gridkey_kwargs["delimiters"] 
    labels_fontsize = gridkey_kwargs.get('labels_fontsize')
    fontsize = gridkey_kwargs.get('fontsize')

    # Auto parser for gridkey - implemented by SangyuXu
    if gridkey == "auto" or gridkey == True:
        if experiment_label is not None:
            gridkey = list(np.concatenate([experiment_label, x1_level]))
        else:
            temp_groups = ";".join(all_plot_groups)
            for delimiter in gridkey_delimiters:
                temp_groups = temp_groups.replace(delimiter, ";")
            temp_groups = [i.strip() for i in temp_groups.split(';')]
            unique_groups = list(set(temp_groups))
            rank = [sum([temp_groups.index(i) for i in temp_groups if(j in i)]) for j in unique_groups]
            gridkey = [x for _,x in sorted(zip(rank,unique_groups))]
    
    # Raise error if there are more than 2 items in any idx and gridkey_merge_pairs is True and is_paired is not None
    if gridkey_merge_pairs and is_paired is not None:
        for i in idx:
            if len(i) > 2:
                warnings.warn(
                    "gridkey_merge_pairs=True only works if all idx in tuples have only two items. gridkey_merge_pairs has automatically been set to False"
                )
                gridkey_merge_pairs = False
                break
    elif gridkey_merge_pairs and is_paired is None:
        warnings.warn(
            "gridkey_merge_pairs=True is only applicable for paired data."
        )
        gridkey_merge_pairs = False

    # Checks for gridkey_merge_pairs and is_paired; if both are true, "merges" the gridkey per pair
    if gridkey_merge_pairs and is_paired is not None:
        groups_for_gridkey = []
        for i in idx:
            groups_for_gridkey.append(i[1])
    else:
        groups_for_gridkey = all_plot_groups

    # raise errors if gridkey is not a list, or if the list is empty
    if isinstance(gridkey, list) is False:
        raise TypeError("gridkey must be a list (or a string 'auto').")
    if any(isinstance(i, str) is False for i in gridkey):
        raise TypeError("gridkey must contain only strings.")
    if len(gridkey) == 0:
        warnings.warn("gridkey is an empty list.")

    # raise Warning if an item in gridkey is not contained in any idx
    for i in gridkey:
        in_idx = 0
        for j in groups_for_gridkey:
            if i in j:
                in_idx += 1
        if in_idx == 0:
            if is_paired is not None:
                warnings.warn(
                    i
                    + " is not in any idx. Please check. Alternatively, merging gridkey pairs may not be suitable for your data; try passing gridkey_merge_pairs=False."
                )
            else:
                warnings.warn(i + " is not in any idx. Please check.")

    # Populate table: checks if idx for each column contains rowlabel name
    # IF so, marks that element as present w black dot (default "\u25CF"), or space if not present
    table_cellcols = []
    for i in gridkey:
        thisrow = []
        for q in groups_for_gridkey:
            if str(i) in q:
                thisrow.append(gridkey_marker)
            else:
                thisrow.append("")
        table_cellcols.append(thisrow)

    # Adds a row for Ns with the Ns values
    if gridkey_show_Ns:
        gridkey.append("Ns")
        list_of_Ns = []
        for i in groups_for_gridkey:
            list_of_Ns.append(str(plot_data.groupby(xvar, observed=False).count()[yvar].loc[i]))
        table_cellcols.append(list_of_Ns)

    # Adds a row for effectsizes with effectsize values
    if gridkey_show_es and not horizontal:
        gridkey.append("\u0394")
        effsize_list = []
        results_list = results.test.to_list()

        # get the effect size, append + or -, 2 dec places
        for i in enumerate(groups_for_gridkey):
            if i[1] in results_list:
                curr_esval = results.loc[results["test"] == i[1]]["difference"].iloc[0]
                curr_esval_str = np.format_float_positional(
                    curr_esval,
                    precision=2,
                    sign=True,
                    trim="k",
                    min_digits=2,
                )
                effsize_list.append(curr_esval_str)
            else:
                effsize_list.append("-")

        table_cellcols.append(effsize_list)

    # Set the axes to plot on
    if float_contrast or horizontal:
        ax_to_plot = rawdata_axes
    else:
        ax_to_plot = contrast_axes

    # Add delta-delta or mini_meta details to the table
    if show_mini_meta or show_delta2:
        if show_delta2:
            added_group_name = ["Deltas' g"] if effect_size == "hedges_g" else ["Delta-Delta"]
        else:
            added_group_name = ["Weighted Delta"]
        gridkey = added_group_name + gridkey
        table_cellcols = [[""]*len(table_cellcols[0])] + table_cellcols

        if not horizontal and show_delta2:
            extra_table_cellcols = [[] for i in range(len(table_cellcols))]

            for group_idx, group_vals in enumerate(extra_table_cellcols):
                if group_idx == 0:
                    added_group = [gridkey_marker]
                elif gridkey_show_es and (group_idx == len(extra_table_cellcols)-1) and not horizontal:
                    added_delta_effectsize = delta_delta.difference
                    added_delta_effectsize_str = np.format_float_positional(
                                                                            added_delta_effectsize,
                                                                            precision=2,
                                                                            sign=True,
                                                                            trim="k",
                                                                            min_digits=2,
                                                                        )
                    added_group = [added_delta_effectsize_str]
                else:
                    added_group = ['']
                for n in added_group:
                    group_vals.append(n)

        elif horizontal or show_mini_meta:
            for group_idx, group_vals in enumerate(table_cellcols):
                if group_idx == 0:
                    added_group = [gridkey_marker]
                elif gridkey_show_es and (group_idx == len(table_cellcols)-1) and not horizontal:
                    added_delta_effectsize = delta_delta.difference if show_delta2 else mini_meta.difference
                    added_delta_effectsize_str = np.format_float_positional(
                                                                            added_delta_effectsize,
                                                                            precision=2,
                                                                            sign=True,
                                                                            trim="k",
                                                                            min_digits=2,
                                                                        )
                    added_group = [added_delta_effectsize_str]
                else:
                    added_group = ['']
                for n in added_group:
                    group_vals.append(n)

    # Create the table object
    def add_table(celltext, bbox, rowlabels=None):
        gridkey_to_plot = ax_to_plot.table(
                    cellText=celltext,
                    rowLabels=rowlabels,
                    cellLoc="center",
                    bbox=bbox,
        )
        return gridkey_to_plot

    if horizontal:
        # Convert the cells format for horizontal table plotting
        converted_list = []
        for j in range(0, len(table_cellcols[0])):
            temp_list = []
            for i in table_cellcols:
                temp_list.append(i[j])
            converted_list.append(temp_list)

        gridkey_to_plot = add_table(celltext = converted_list, bbox = [-len(gridkey) * 0.2, 0, len(gridkey) * 0.2, 1])

        # Add the column labels as text below the table
        text_locs = np.arange((-len(gridkey)*0.2) +0.1, 0, 0.2)
        for loc, txt in zip(text_locs, gridkey):
            ax_to_plot.text(
                        loc+0.04, 
                        -0.01, 
                        txt, 
                        transform=ax_to_plot.transAxes, 
                        ha='right',
                        rotation=45,
                        fontsize=labels_fontsize if labels_fontsize is not None else 10,
                        va='top',
                    )
    else:
        # Plot the table for vertical format
        if show_mini_meta:
            gridkey_to_plot = add_table(celltext = table_cellcols, rowlabels=gridkey, bbox = [0, -len(gridkey) * 0.1 - 0.05, 1, len(gridkey) * 0.1])
        elif show_delta2:
            gridkey_to_plot = add_table(celltext = table_cellcols, rowlabels=gridkey, bbox = [0, -len(gridkey) * 0.1 - 0.05, 0.75, len(gridkey) * 0.1])
            extra_gridkey  =  add_table(celltext = extra_table_cellcols, bbox = [0.78, -len(gridkey) * 0.1 - 0.05, 0.15, len(gridkey) * 0.1])
        else:
            gridkey_to_plot = add_table(celltext = table_cellcols, rowlabels=gridkey, bbox = [0, -len(gridkey) * 0.1 - 0.05, 1, len(gridkey) * 0.1]) 

        # modifies row label cells
        for cell in gridkey_to_plot._cells:
            if cell[1] == -1:
                gridkey_to_plot._cells[cell].visible_edges = "open"
                gridkey_to_plot._cells[cell].set_text_props(**{"ha": "right"})

    if fontsize is not None:
        gridkey_to_plot.auto_set_font_size(False)
        gridkey_to_plot.set_fontsize(fontsize)
        if show_delta2 and not horizontal:
            extra_gridkey.auto_set_font_size(False)
            extra_gridkey.set_fontsize(fontsize)

    if labels_fontsize is not None and not horizontal:
        gridkey_to_plot.auto_set_font_size(False)
        for cell in gridkey_to_plot._cells:
            if cell[1] == -1:
                gridkey_to_plot._cells[cell].set_text_props(**{"fontsize": labels_fontsize})

    # turns off both x axes
    if horizontal:
        rawdata_axes.get_yaxis().set_visible(False)
        contrast_axes.get_yaxis().set_visible(False)
    else:
        rawdata_axes.get_xaxis().set_visible(False)
        contrast_axes.get_xaxis().set_visible(False)

def barplotter(
        xvar: str, 
        yvar: str, 
        all_plot_groups: list, 
        rawdata_axes: axes.Axes, 
        plot_data: pd.DataFrame, 
        raw_colors: str, 
        plot_palette_raw: dict, 
        color_col: str,
        barplot_kwargs: dict, 
        horizontal: bool
    ):
    """
    Add bars to the raw data plot.

    Parameters
    ----------
    xvar : str
        Column name of the x variable.
    yvar : str
        Column name of the y variable.
    all_plot_groups : list
        List of all plot groups.
    rawdata_axes : object
        Matplotlib axis object to plot on.
    plot_data : object (Dataframe)
        Dataframe of the plot data.
    raw_colors : str
        Color of the bar.
    plot_palette_raw : dict
        Dictionary of colors used in the bar plot.
    color_col : str
        Column name of the color column.
    barplot_kwargs : dict
        Keyword arguments for the barplot.
    horizontal : bool
        If the plot is horizontal.
    """
    # Check if the custom_palette is a dictionary with two keys 0 and 1 (for filled bar coloring)
    filled_bars = True if len(plot_palette_raw.keys())==2 and all(k in plot_palette_raw for k in [1, 0]) else False

    bar_width = barplot_kwargs.get('width', 0.5)
    fontsize = barplot_kwargs.pop('fontsize')

    x_label, y_label = rawdata_axes.get_xlabel(), rawdata_axes.get_ylabel()
    if horizontal:
        x_var, y_var, orient = np.ones(len(all_plot_groups)), all_plot_groups, "h"
    else:
        x_var, y_var, orient = all_plot_groups, np.ones(len(all_plot_groups)), "v"

    # Create bar1_df with basic columns
    bar1_df = pd.DataFrame({
        xvar: x_var, 
        "proportion": y_var
    })

    # Handle colors
    if color_col:
        # Get first color value for each group
        color_mapping = plot_data.groupby(xvar, observed=False)[color_col].first()
        bar1_df[color_col] = [color_mapping.get(group) for group in all_plot_groups]
        
        # Map colors, defaulting to bar_color if no match
        edge_colors = [
            plot_palette_raw.get(hue_val, raw_colors) 
            for hue_val in bar1_df[color_col]
        ]
    else:
        edge_colors = len(all_plot_groups)*['black',] if filled_bars else raw_colors

    bar1 = sns.barplot(
        data=bar1_df,
        x=xvar,
        y="proportion",
        ax=rawdata_axes,
        order=all_plot_groups,
        linewidth=1 if filled_bars else 2,
        facecolor=plot_palette_raw[0] if filled_bars else (1, 1, 1, 0),
        edgecolor=edge_colors,
        zorder=1,
        orient=orient,
    )

    if filled_bars:
        barplot_kwargs['facecolor'] = plot_palette_raw[1]
        barplot_kwargs['edgecolor'] = 'black'
        barplot_kwargs['linewidth'] = 1
    else:
        barplot_kwargs['palette'] = plot_palette_raw

    bar2 = sns.barplot(
        data=plot_data,
        x=yvar if horizontal else xvar,
        y=xvar if horizontal else yvar,
        hue=xvar if color_col is None else color_col,
        ax=rawdata_axes,
        order=all_plot_groups,
        dodge=False,
        zorder=1,
        orient=orient,
        **barplot_kwargs
    )

    # adjust the width of bars
    if horizontal:
        for bar in bar1.patches:
            y = bar.get_y()
            height = bar.get_height()
            centre = y + height / 2.0
            bar.set_y(centre - bar_width / 2.0)
            bar.set_height(bar_width)
    else:
        for bar in bar1.patches:
            x = bar.get_x()
            width = bar.get_width()
            centre = x + width / 2.0
            bar.set_x(centre - bar_width / 2.0)
            bar.set_width(bar_width)

    # reset the x and y labels
    rawdata_axes.set_xlabel(x_label)
    rawdata_axes.set_ylabel(y_label)

    if horizontal:
        rawdata_axes.set_yticks(rawdata_axes.get_yticks())
        rawdata_axes.set_yticklabels(rawdata_axes.get_yticklabels(), fontsize = fontsize)
    else:
        rawdata_axes.set_xticks(rawdata_axes.get_xticks())
        rawdata_axes.set_xticklabels(rawdata_axes.get_xticklabels(), fontsize = fontsize)

def table_for_horizontal_plots(
        effectsize_df: object, 
        ax: axes.Axes, 
        contrast_axes: axes.Axes, 
        ticks_to_plot: list, 
        show_mini_meta: bool, 
        show_delta2: bool, 
        table_kwargs: dict,
        ticks_to_skip: list
    ):
    """
    Add table axes for showing the deltas for horizontal plots.

    Parameters
    ----------
    effectsize_df : object
        Effect size DABEST object.
    ax : object
        Matplotlib axis object to plot the table axes.
    contrast_axes : object
        Matplotlib axis object to plot the contrast axes.
    ticks_to_plot : list
        List of indices of the contrast objects.
    show_mini_meta : bool
        Whether to show the mini meta-analysis.
    show_delta2 : bool
        Whether to show the delta-delta.
    table_kwargs : dict
        Keyword arguments for the table.
    ticks_to_skip:  list
        List of ticks to skip in the table.
    """

    table_color = table_kwargs['color']
    table_alpha = table_kwargs['alpha']
    table_font_size = table_kwargs['fontsize']
    table_text_color = table_kwargs['text_color']
    text_units = table_kwargs['text_units']
    table_font_size -= 2 if text_units != '' else 0
    control_marker = table_kwargs['control_marker']  
    fontsize_label = table_kwargs['fontsize_label']
    label = table_kwargs['label']

    ### Create a table of deltas
    cols=['Δ','N']
    lst = []
    for n in np.arange(0, len(effectsize_df.results.difference), 1):
        lst.append([effectsize_df.results.difference[n], 0])
    if show_mini_meta:
        lst.append([effectsize_df.mini_meta.difference, 0])
    elif show_delta2:
        lst.append([effectsize_df.delta_delta.difference, 0])
    tab = pd.DataFrame(lst, columns=cols)

    ### Plot the text
    if show_mini_meta or show_delta2:
        new_ticks = ticks_to_plot + [max(ticks_to_plot)+1]
    else:
        new_ticks = ticks_to_plot.copy()
    for i,loc in zip(tab.index, new_ticks):
        ax.text(0.5, loc, "{:+.2f}".format(tab.iloc[i,0])+text_units, ha="center", va="center", color=table_text_color, size=table_font_size)

    # Plot the dashes
    if control_marker is not None:
        for loc in ticks_to_skip:
            ax.text(0.5, loc, control_marker, ha="center", va="center", color=table_text_color, size=table_font_size)

    ### Parameters for table
    ax.axvspan(0, 1, facecolor=table_color, alpha=table_alpha)  #### Plot the background color
    ax.set_xticks([0.5])
    ax.set_xticklabels([])
    ax.set_ylim(contrast_axes.get_ylim())
    ax.set_yticks([])
    ax.set_yticklabels([])
    ax.tick_params(left=False, bottom=False)
    ax.set_xlabel(label, fontsize=fontsize_label) # Set the x-axis label - hardcoded for now
    sns.despine(ax=ax, left=True, bottom=True)


def add_counts_to_prop_plots(
        plot_data: pd.DataFrame, 
        xvar: str, 
        yvar: str, 
        rawdata_axes: axes.Axes, 
        horizontal: bool, 
        is_paired: bool, 
        prop_sample_counts_kwargs: dict
    ):
    """
    Add counts to the proportion plots.

    Parameters
    ----------
    plot_data : object (Dataframe)
        Dataframe of the plot data.
    xvar : str
        Column name of the x variable.
    yvar : str
        Column name of the y variable.
    rawdata_axes : axes.Axes
        Matplotlib axis object to plot on.
    horizontal : bool
        If the plot is horizontal.
    is_paired : bool
        Whether the data is paired.
    prop_sample_counts_kwargs : dict
        Keyword arguments for the sample counts.
    """

    # Group orders
    if isinstance(plot_data[xvar].dtype, pd.CategoricalDtype):
        sample_size_text_order = pd.unique(plot_data[xvar]).categories
    else:
        sample_size_text_order = pd.unique(plot_data[xvar])

    # Get the sample size values
    ones, zeros = plot_data[plot_data[yvar] == 1], plot_data[plot_data[yvar] == 0]

    sample_size_val1 = ones.groupby(xvar, observed=False)[yvar].count().reindex(index=sample_size_text_order)
    sample_size_val0 = zeros.groupby(xvar, observed=False)[yvar].count().reindex(index=sample_size_text_order)

    if "fontsize" not in prop_sample_counts_kwargs.keys():
        fontsize = 8 if horizontal else 10
        fontsize -= 2 if is_paired else 0
        prop_sample_counts_kwargs.update({'fontsize': fontsize})

    for sample_text_x, sample_text_y0, sample_text_y1 in zip(
                                                            np.arange(0, len(sample_size_text_order) + 1, 1), 
                                                            sample_size_val0,
                                                            sample_size_val1,
                                                            ):
        if horizontal:
            rawdata_axes.text(0.05, sample_text_x, sample_text_y1, **prop_sample_counts_kwargs)
            rawdata_axes.text(0.95, sample_text_x, sample_text_y0, **prop_sample_counts_kwargs)
        else:
            rawdata_axes.text(sample_text_x, 0.05, sample_text_y1, **prop_sample_counts_kwargs)
            rawdata_axes.text(sample_text_x, 0.95, sample_text_y0, **prop_sample_counts_kwargs)

# %% ../nbs/API/plot_tools.ipynb 6
def swarmplot(
    data: pd.DataFrame,
    x: str,
    y: str,
    ax: axes.Axes,
    order: List = None,
    hue: str = None,
    palette: Union[Iterable, str] = "black",
    zorder: float = 1,
    size: float = 5,
    side: str = "center",
    jitter: float = 1,
    filled: Union[bool, List, Tuple] = True,
    is_drop_gutter: bool = True,
    gutter_limit: float = 0.5,
    horizontal: bool = False,
    **kwargs,
):
    """
    API to plot a swarm plot.

    Parameters
    ----------
    data : pd.DataFrame
        The input data as a pandas DataFrame.
    x : str
        The column in the DataFrame to be used as the x-axis.
    y : str
        The column in the DataFrame to be used as the y-axis.
    ax : axes.Axes
        Matplotlib axes.Axes object for which the plot would be drawn on. Default is None.
    order : List
        The order in which x-axis categories should be displayed. Default is None.
    hue : str
        The column in the DataFrame that determines the grouping for color.
        If None (by default), it assumes that it is being grouped by x.
    palette : Union[Iterable, str]
        The color palette to be used for plotting. Default is "black".
    zorder : int | float
        The z-order for drawing the swarm plot wrt other matplotlib drawings. Default is 1.
    dot_size : int | float
        The size of the markers in the swarm plot. Default is 20.
    side : str
        The side on which points are swarmed ("center", "left", or "right"). Default is "center".
    jitter : int | float
        Determines the distance between points. Default is 1.
    filled : bool | List | Tuple
        Determines whether the dots in the swarmplot are filled or not. If set to False,
        dots are not filled. If provided as a List or Tuple, it should contain boolean values,
        each corresponding to a swarm group in order, indicating whether the dot should be
        filled or not.
    is_drop_gutter : bool
        If True, drop points that hit the gutters; otherwise, readjust them.
    gutter_limit : int | float
        The limit for points hitting the gutters.
    horizontal : bool
        If True, the swarm plot is drawn horizontally. Default is False.
    **kwargs:
        Additional keyword arguments to be passed to the swarm plot.

    Returns
    -------
    axes.Axes
        Matplotlib axes.Axes object for which the swarm plot has been drawn on.
    """
    s = SwarmPlot(data, x, y, ax, order, hue, palette, zorder, size, side, jitter, horizontal)
    ax = s.plot(is_drop_gutter, gutter_limit, ax, filled, horizontal, **kwargs)
    return ax


class SwarmPlot:
    def __init__(
        self,
        data: pd.DataFrame,
        x: str,
        y: str,
        ax: axes.Axes,
        order: List = None,
        hue: str = None,
        palette: Union[Iterable, str] = "black",
        zorder: float = 1,
        size: float = 5,
        side: str = "center",
        jitter: float = 1,
        horizontal: bool = False,
    ):
        """
        Initialize a SwarmPlot instance.

        Parameters
        ----------
        data : pd.DataFrame
            The input data as a pandas DataFrame.
        x : str
            The column in the DataFrame to be used as the x-axis.
        y : str
            The column in the DataFrame to be used as the y-axis.
        ax : axes.Axes
            Matplotlib axes.Axes object for which the plot would be drawn on.
        order : List
            The order in which x-axis categories should be displayed. Default is None.
        hue : str
            The column in the DataFrame that determines the grouping for color.
            If None (by default), it assumes that it is being grouped by x.
        palette : Union[Iterable, str]
            The color palette to be used for plotting. Default is "black".
        zorder : int | float
            The z-order for drawing the swarm plot wrt other matplotlib drawings. Default is 1.
        dot_size : int | float
            The size of the markers in the swarm plot. Default is 20.
        side : str
            The side on which points are swarmed ("center", "left", or "right"). Default is "center".
        jitter : int | float
            Determines the distance between points. Default is 1.
        horizontal : bool
            If True, the swarm plot is drawn horizontally. Default is False.

        Returns
        -------
        None
        """
        self.__x = x
        self.__y = y
        self.__order = order
        self.__hue = hue
        self.__zorder = zorder
        self.__palette = palette
        self.__jitter = jitter

        # Input validation
        self._check_errors(data, ax, size, side)

        self.__size = size * 4
        self.__side = side.lower()
        self.__data = data
        self.__color_col = self.__x if self.__hue is None else self.__hue

        # Generate default values
        if order is None:
            self.__order = self._generate_order()

        # Reformatting
        if not isinstance(self.__palette, dict):
            self.__palette = self._format_palette(self.__palette)
        data_copy = data.copy(deep=True)
        if not isinstance(self.__data[self.__x].dtype, pd.CategoricalDtype):
            # make x column into CategoricalDType to sort by
            data_copy[self.__x] = data_copy[self.__x].astype(
                CategoricalDtype(categories=self.__order, ordered=True)
            )
        data_copy.sort_values(by=[self.__x, self.__y], inplace=True)
        self.__data_copy = data_copy

        x_vals = range(len(self.__order))
        y_vals = self.__data_copy[self.__y]

        x_min = min(x_vals)
        x_max = max(x_vals)
        y_range = max(y_vals) - min(y_vals)
        y_min = min(y_vals) - 0.05 * y_range
        y_max = max(y_vals) + 0.05 * y_range

        if horizontal:
            ax.set_ylim(bottom=x_min - 0.5, top=x_max + 0.5)
            # ylim is set manually to override Axes.autoscale if it hasn't already been scaled at least once
            if ax.get_autoscalex_on():
                ax.set_xlim(left=y_min, right=y_max)
        else:
            ax.set_xlim(left=x_min - 0.5, right=x_max + 0.5)
            # ylim is set manually to override Axes.autoscale if it hasn't already been scaled at least once
            if ax.get_autoscaley_on():
                ax.set_ylim(bottom=y_min, top=y_max)

        figw, figh = ax.get_figure().get_size_inches()
        w = (ax.get_position().xmax - ax.get_position().xmin) * figw
        h = (ax.get_position().ymax - ax.get_position().ymin) * figh
        ax_xspan = ax.get_xlim()[1] - ax.get_xlim()[0]
        ax_yspan = ax.get_ylim()[1] - ax.get_ylim()[0]
        if horizontal:
            ax_xspan, ax_yspan = ax_yspan, ax_xspan

        # increases jitter distance based on number of swarms that is going to be drawn
        jitter = jitter * (1 + 0.05 * (math.log(ax_xspan)))

        gsize = (
            math.sqrt(self.__size) * 1.0 / (70 / jitter) * ax_xspan * 1.0 / (w * 0.8)
        )
        dsize = (
            math.sqrt(self.__size) * 1.0 / (70 / jitter) * ax_yspan * 1.0 / (h * 0.8)
        )
        self.__gsize = gsize
        self.__dsize = dsize

    def _check_errors(
        self, data: pd.DataFrame, ax: axes.Axes, size: float, side: str
    ) -> None:
        """
        Check the validity of input parameters. Raises exceptions if detected.

        Parameters
        ----------
        data : pd.Dataframe
            Input data used for generation of the swarmplot.
        ax : axes.Axes
            Matplotlib axes.Axes object for which the plot would be drawn on.
        size : int | float
            scalar value determining size of dots of the swarmplot.
        side: str
            The side on which points are swarmed ("center", "left", or "right"). Default is "center".

        Returns
        -------
        None
        """
        # Type enforcement
        if not isinstance(data, pd.DataFrame):
            raise ValueError("`data` must be a Pandas Dataframe.")
        if not isinstance(ax, axes.Axes):
            raise ValueError(
                f"`ax` must be a Matplotlib axes.Axes. The current `ax` is a {type(ax)}"
            )
        if not isinstance(size, (int, float)):
            raise ValueError("`size` must be a scalar or float.")
        if not isinstance(side, str):
            raise ValueError(
                "Invalid `side`. Must be one of 'center', 'right', or 'left'."
            )
        if not isinstance(self.__x, str):
            raise ValueError("`x` must be a string.")
        if not isinstance(self.__y, str):
            raise ValueError("`y` must be a string.")
        if not isinstance(self.__zorder, (int, float)):
            raise ValueError("`zorder` must be a scalar or float.")
        if not isinstance(self.__jitter, (int, float)):
            raise ValueError("`jitter` must be a scalar or float.")
        if not isinstance(self.__palette, (str, Iterable)):
            raise ValueError(
                "`palette` must be either a string indicating a color name or an Iterable."
            )
        if self.__hue is not None and not isinstance(self.__hue, str):
            raise ValueError("`hue` must be either a string or None.")
        if self.__order is not None and not isinstance(self.__order, Iterable):
            raise ValueError("`order` must be either an Iterable or None.")

        # More thorough input validation checks
        if self.__x not in data.columns:
            err = "{0} is not a column in `data`.".format(self.__x)
            raise IndexError(err)
        if self.__y not in data.columns:
            err = "{0} is not a column in `data`.".format(self.__y)
            raise IndexError(err)
        if self.__hue is not None and self.__hue not in data.columns:
            err = "{0} is not a column in `data`.".format(self.__hue)
            raise IndexError(err)

        color_col = self.__x if self.__hue is None else self.__hue
        if self.__order is not None:
            for group_i in self.__order:
                if group_i not in pd.unique(data[self.__x]):
                    err = "{0} in `order` is not in the '{1}' column of `data`.".format(
                        group_i, self.__x
                    )
                    raise IndexError(err)

        if isinstance(self.__palette, str) and self.__palette.strip() == "":
            err = "`palette` cannot be an empty string. It must be either a string indicating a color name or an Iterable."
            raise ValueError(err)
        if isinstance(self.__palette, dict):
            for group_i, color_i in self.__palette.items():
                if group_i not in pd.unique(data[color_col]):
                    err = (
                        "{0} in `palette` is not in the '{1}' column of `data`.".format(
                            group_i, color_col
                        )
                    )
                    raise IndexError(err)
                if isinstance(color_i, str) and color_i.strip() == "":
                    err = "The color mapping for {0} in `palette` is an empty string. It must contain a color name.".format(
                        group_i
                    )
                    raise ValueError(err)

        if side.lower() not in ["center", "right", "left"]:
            raise ValueError(
                "Invalid `side`. Must be one of 'center', 'right', or 'left'."
            )

        return None

    def _generate_order(self) -> List:
        """
        Generates order value that determines the order in which x-axis categories should be displayed.

        Parameters
        ----------
        None

        Returns
        -------
        List:
            contains the order in which the x-axis categories should be displayed.
        """
        if isinstance(self.__data[self.__x].dtype, pd.CategoricalDtype):
            order = pd.unique(self.__data[self.__x]).categories.tolist()
        else:
            order = pd.unique(self.__data[self.__x]).tolist()

        return order

    def _format_palette(self, palette: Union[str, List, Tuple]) -> Dict:
        """
        Reformats palette into appropriate Dictionary form for swarm plot

        Parameters
        ----------
        palette: str | List | Tuple
            The color palette used for the swarm plot. Conventions are based on Matplotlib color
            specifications.

            Could be a singular string value - in which case, would be a singular color name.
            In the case of a List or Tuple - it could be a Sequence of color names or RGB(A) values.

        Returns
        -------
        Dict:
            Dictionary mapping unique groupings in the color column (of the data used for the swarm plot)
            to a color name (str) or a RGB(A) value (Tuple[float, float, float] | List[float, float, float]).
        """
        reformatted_palette = dict()
        groups = pd.unique(self.__data[self.__color_col]).tolist()

        if isinstance(palette, str):
            for group_i in groups:
                reformatted_palette[group_i] = palette
        if isinstance(palette, (list, tuple)):
            if len(groups) != len(palette):
                err = (
                    "unique values in '{0}' column in `data` "
                    "and `palette` do not have the same length. Number of unique values is {1} "
                    "while length of palette is {2}. The assignment of the colors in the "
                    "palette will be cycled."
                ).format(self.__color_col, len(groups), len(palette))
                warnings.warn(err)
            for i, group_i in enumerate(groups):
                reformatted_palette[group_i] = palette[i % len(palette)]

        return reformatted_palette

    def _swarm(
        self, values: Iterable[float], gsize: float, dsize: float, side: str
    ) -> pd.Series:
        """
        Perform the swarm algorithm to position points without overlap.

        Parameters
        ----------
        values : Iterable[int | float]
            The values to be plotted.
        gsize : int | float
            The size of the gap between points.
        dsize : int | float
            The size of the markers.
        side : str
            The side on which points are swarmed ("center", "left", or "right").

        Returns
        -------
        pd.Series:
            The x-offset values for the swarm plot.
        """
        # Input validation
        if not isinstance(values, Iterable):
            raise ValueError("`values` must be an Iterable")
        if not isinstance(gsize, (int, float)):
            raise ValueError("`gsize` must be a scalar or float.")
        if not isinstance(dsize, (int, float)):
            raise ValueError("`dsize` must be a scalar or float.")

        # Sorting algorithm based off of: https://github.com/mgymrek/pybeeswarm
        points_data = pd.DataFrame({
            "y": [yval * 1.0 / dsize for yval in values],
            "x": np.zeros(len(values), dtype=float)  # Initialize with float zeros
        })
        for i in range(1, points_data.shape[0]):
            y_i = points_data["y"].values[i]
            points_placed = points_data[0:i]
            is_points_overlap = (
                abs(y_i - points_placed["y"]) < 1
            )  # Checks if y_i is overlapping with any points already placed
            if any(is_points_overlap):
                points_placed = points_placed[is_points_overlap]
                x_offsets = points_placed["y"].apply(
                    lambda y_j: math.sqrt(1 - (y_i - y_j) ** 2)
                )
                if side == "center":
                    potential_x_offsets = pd.Series(
                        [0]
                        + (points_placed["x"] + x_offsets).tolist()
                        + (points_placed["x"] - x_offsets).tolist()
                    )
                if side == "right":
                    potential_x_offsets = pd.Series(
                        [0] + (points_placed["x"] + x_offsets).tolist()
                    )
                if side == "left":
                    potential_x_offsets = pd.Series(
                        [0] + (points_placed["x"] - x_offsets).tolist()
                    )
                bad_x_offsets = []
                for x_i in potential_x_offsets:
                    dists = (y_i - points_placed["y"]) ** 2 + (
                        x_i - points_placed["x"]
                    ) ** 2
                    if any([item < 0.999 for item in dists]):
                        bad_x_offsets.append(True)
                    else:
                        bad_x_offsets.append(False)
                potential_x_offsets[bad_x_offsets] = np.inf
                abs_potential_x_offsets = [abs(_) for _ in potential_x_offsets]
                valid_x_offset = potential_x_offsets[
                    abs_potential_x_offsets.index(min(abs_potential_x_offsets))
                ]
                points_data.loc[i, "x"] = valid_x_offset
            else:
                points_data.loc[i, "x"] = 0

        points_data.loc[np.isnan(points_data["y"]), "x"] = np.nan

        return points_data["x"] * gsize

    def _adjust_gutter_points(
        self,
        points_data: pd.DataFrame,
        x_position: float,
        is_drop_gutter: bool,
        gutter_limit: float,
        value_column: str,
    ) -> pd.DataFrame:
        """
        Adjust points that hit the gutters or drop them based on the provided conditions.

        Parameters
        ----------
        points_data: pd.DataFrame
            Data containing coordinates of points for the swarm plot.
        x_position: int | float
            X-coordinate of the center of a singular swarm group of the swarm plot
        is_drop_gutter : bool
            If True, drop points that hit the gutters; otherwise, readjust them.
        gutter_limit : int | float
            The limit for points hitting the gutters.
        value_column : str
            column in points_data that contains the coordinates for the points in the axis against the gutter

        Returns
        -------
        pd.DataFrame:
            DataFrame with adjusted points based on the gutter limit.
        """
        if self.__side == "center":
            gutter_limit = gutter_limit / 2

        hit_gutter = abs(points_data[value_column] - x_position) >= gutter_limit
        total_num_of_points = points_data.shape[0]
        num_of_points_hit_gutter = points_data[hit_gutter].shape[0]
        if any(hit_gutter):
            if is_drop_gutter:
                # Drop points that hit gutter
                points_data.drop(points_data[hit_gutter].index.to_list(), inplace=True)
                err = (
                    "{0:.1%} of the points cannot be placed. "
                    "You might want to decrease the size of the markers."
                ).format(num_of_points_hit_gutter / total_num_of_points)
                warnings.warn(err)
            else:
                for i in points_data[hit_gutter].index:
                    points_data.loc[i, value_column] = np.sign(
                        points_data.loc[i, value_column]
                    ) * (x_position + gutter_limit)

        return points_data

    def plot(
        self,
        is_drop_gutter: bool,
        gutter_limit: float,
        ax: axes.Subplot,
        filled: Union[bool, List, Tuple],
        horizontal: bool,
        **kwargs,
    ) -> axes.Axes:
        """
        Generate a swarm plot.

        Parameters
        ----------
        is_drop_gutter : bool
            If True, drop points that hit the gutters; otherwise, readjust them.
        gutter_limit : int | float
            The limit for points hitting the gutters.
        ax : axes.Axes
            The matplotlib figure object to which the swarm plot will be added.
        filled : bool | List | Tuple
            Determines whether the dots in the swarmplot are filled or not. If set to False,
            dots are not filled. If provided as a List or Tuple, it should contain boolean values,
            each corresponding to a swarm group in order, indicating whether the dot should be
            filled or not.
        **kwargs:
            Additional keyword arguments to be passed to the scatter plot.

        Returns
        -------
        axes.Axes:
            The matplotlib axes containing the swarm plot.
        """
        # Input validation
        if not isinstance(is_drop_gutter, bool):
            raise ValueError("`is_drop_gutter` must be a boolean.")
        if not isinstance(gutter_limit, (int, float)):
            raise ValueError("`gutter_limit` must be a scalar or float.")
        if not isinstance(filled, (bool, list, tuple)):
            raise ValueError("`filled` must be a boolean, list or tuple.")
        
        fontsize = kwargs.pop('fontsize', 12)

        # More thorough input validation checks
        if isinstance(filled, (list, tuple)):
            if len(filled) != len(self.__order):
                err = (
                    "There are {0} unique values in `x` column in `data` "
                    "but `filled` has a length of {1}. If `filled` is a list "
                    "or a tuple, it must have the same length as the number of "
                    "unique values/groups in the `x` column of data."
                ).format(len(self.__order), len(filled))
                raise ValueError(err)
            if not all(isinstance(_, bool) for _ in filled):
                raise ValueError("All values in `filled` must be a boolean.")

        # Assumptions are that self.__data_copy is already sorted according to self.__order
        x_position = (
            0  # x-coordinate of center of each individual swarm of the swarm plot
        )
        x_tick_tabels = []
        for group_i, values_i in self.__data_copy.groupby(self.__x, observed=False):
            x_new = []
            values_i_y = values_i[self.__y]
            x_offset = self._swarm(
                values=values_i_y,
                gsize=self.__gsize,
                dsize=self.__dsize,
                side=self.__side,
            )
            x_new = [
                x_position + offset for offset in x_offset
            ]  # apply x-offsets based on _swarm algo
            values_i["x_new"] = x_new
            values_i = self._adjust_gutter_points(
                values_i, x_position, is_drop_gutter, gutter_limit, "x_new"
            )
            x_tick_tabels.extend([group_i])
            x_position = x_position + 1

            if values_i.empty:
                ax.scatter(
                    values_i["x_new"] if not horizontal else values_i[self.__y],
                    values_i[self.__y] if not horizontal else values_i["x_new"],
                    s=self.__size,
                    zorder=self.__zorder,
                    **kwargs,
                )
                continue

            if self.__hue is not None:
                # color swarms based on `hue` column
                cmap_values, index = np.unique(
                    values_i[self.__hue], return_inverse=True
                )
                cmap = []
                for cmap_group_i in cmap_values:
                    cmap.append(self.__palette[cmap_group_i])

                cmap = ListedColormap(cmap)
                ax.scatter(
                    values_i["x_new"] if not horizontal else values_i[self.__y],
                    values_i[self.__y] if not horizontal else values_i["x_new"],
                    s=self.__size,
                    c=index,
                    cmap=cmap,
                    zorder=self.__zorder,
                    edgecolor="face",
                    **kwargs,
                )

            else:
                # color swarms based on `x` column
                if not isinstance(filled, bool):
                    facecolor = (
                        "none"
                        if not filled[x_position - 1]
                        else self.__palette[group_i]
                    )
                else:
                    facecolor = "none" if not filled else self.__palette[group_i]

                ax.scatter(
                    values_i["x_new"] if not horizontal else values_i[self.__y],
                    values_i[self.__y] if not horizontal else values_i["x_new"],
                    s=self.__size,
                    zorder=self.__zorder,
                    facecolor=facecolor,
                    edgecolor=self.__palette[group_i],
                    label=group_i,
                    **kwargs,
                )

        # Handling of legends
        # This is currently a workaround because c and cmap is unable to map the labels when calling scatter()
        # labels has to be used to designate legend labels and handles in scatter() due to the potential calling of ax.get_legend_handles_labels()
        if self.__hue is not None:
            for cmap_group_i in self.__palette:
                ax.scatter(
                    [],
                    [],
                    color=self.__palette[cmap_group_i],
                    label=cmap_group_i,
                )

        if horizontal:
            ax.get_yaxis().set_ticks(np.arange(x_position))
            ax.get_yaxis().set_ticklabels(x_tick_tabels, fontsize = fontsize)
        else:
            ax.get_xaxis().set_ticks(np.arange(x_position))
            ax.get_xaxis().set_ticklabels(x_tick_tabels, fontsize = fontsize)
            
        return ax
