# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.colors as mcolors


def stripplot(
    ax,
    data,
    x_key,
    y_key,
    x_label=None,
    y_label=None,
    title=None,
    units='',
    color_palette='Set2',
):
    """Generates a strip plot with overlaid mean and confidence intervals.

    This function creates a seaborn stripplot to visualize the distribution of
    data points, and overlays a pointplot to show the mean and 95% confidence
    interval for each group. It also annotates the mean value for each group
    on the plot.

    Args:
        ax: The matplotlib axes object to draw the plot on.
        data: A pandas DataFrame containing the data to plot.
        x_key: The name of the column in `data` to group by for the x-axis.
        y_key: The name of the column in `data` for the y-axis values.
        x_label: Optional label for the x-axis. If None, `x_key` is used.
        y_label: Optional label for the y-axis. If None, `y_key` is used.
        title: Optional title for the plot.
        units: Optional string to append to the mean value labels (e.g., 'ms').
        color_palette: The seaborn color palette to use for the plot.

    Usage::

        import pandas as pd
        import matplotlib.pyplot as plt
        import colabutils

        # Sample data
        df = pd.DataFrame({
            'version': ['A', 'A', 'A', 'B', 'B', 'B'],
            'latency': [100, 105, 102, 110, 112, 115]
        })

        # Create plot
        sns.set_theme(style='darkgrid')
        _, axes = plt.subplots(1, 1, figsize=(20, 15))
        colabutils.plot.stripplot(ax=axes,
                        data=df,
                        x_key='version',
                        y_key='latency',
                        x_label='Version',
                        y_label='Latency (ms)',
                        title='Latency Comparison',
                        units='ms',
                        color_palette='Set1')
        plt.show()
    """
    sns.set_palette(color_palette)

    # Use seaborn.stripplot to show individual data points for each group.
    # `hue` is used to give each group a distinct color.
    sns.stripplot(
        x=x_key,
        y=y_key,
        data=data,
        hue=x_key,
        size=8,
        alpha=0.6,
        ax=ax,
        legend=False,
    )

    # Overlay a pointplot to display the mean and 95% confidence interval.
    # `join=False` prevents drawing lines between points of different groups.
    point_plot = sns.pointplot(
        x=x_key,
        y=y_key,
        data=data,
        hue=x_key,
        ax=ax,
        join=False,
        markers='d',
        errorbar=('ci', 95),
        capsize=0.05,
    )

    # Get the unique groups to ensure a consistent order for colors and labels.
    unique_groups = data[x_key].unique()

    # Get the colors used by seaborn for each group to ensure consistent
    # coloring for the mean value labels.
    colors = sns.color_palette(color_palette, n_colors=len(unique_groups))

    # Manually create legend handles with circular markers in the palette color.
    legend_handles = [
        mlines.Line2D(
            [],
            [],
            color=colors[i],
            marker='o',
            linestyle='None',
            markersize=8,
            label=group,
        )
        for i, group in enumerate(unique_groups)
    ]
    ax.legend(handles=legend_handles, title=x_label)

    # Add data labels for the mean values with corresponding colors
    for i, group in enumerate(unique_groups):
        mean_value = data[data[x_key] == group][y_key].mean()
        # Darken the color for the text annotation to improve contrast and
        # readability against the plot background.
        r, g, b, a = mcolors.to_rgba(colors[i], alpha=1.0)
        darker_color = (r * 0.6, g * 0.6, b * 0.6, a)
        # Add a text annotation for the mean. A small horizontal offset is added
        # to prevent the label from overlapping with the point marker.
        ax.text(
            i + 0.1,
            mean_value,
            f'{mean_value:.0f}{units}',
            ha='left',
            va='center',
            color=darker_color,
            fontweight='bold',
        )

    ax.set_ylabel(y_label or y_key)
    ax.set_xlabel(x_label or x_key)
    if title:
        ax.set_title(title)