"""
分组箱线图
"""

import matplotlib.pyplot as plt
import seaborn as sns
from vizpaint.utils import set_chinese_font

import numpy as np
def create_grouped_boxplot(data, x_column, y_column, group_column=None,
                           title="分组箱线图", xlabel="X轴", ylabel="Y轴",
                           figsize=(12, 8), colors=None, legend_loc='upper right',
                           rotation=45):
    """
    创建分组箱线图

    参数:
    ----------
    data : DataFrame
        包含数据的DataFrame
    x_column : str
        X轴列名
    y_column : str
        Y轴列名
    group_column : str, optional
        分组列名
    title : str, optional
        图表标题
    xlabel : str, optional
        X轴标签
    ylabel : str, optional
        Y轴标签
    figsize : tuple, optional
        图表尺寸
    colors : list, optional
        颜色列表
    legend_loc : str, optional
        图例位置
    rotation : int, optional
        X轴标签旋转角度

    返回:
    ----------
    fig, ax : 图表对象和坐标轴对象
    """
    set_chinese_font()
    fig, ax = plt.subplots(figsize=figsize)

    if group_column is None:
        # 简单箱线图
        if colors is None:
            colors = plt.cm.Set3(0.3)

        ax.boxplot([data[data[x_column] == cat][y_column].dropna()
                    for cat in data[x_column].unique()])

        ax.set_xticklabels(data[x_column].unique(), rotation=rotation)

    else:
        # 分组箱线图
        if colors is None:
            colors = plt.cm.Set3(np.linspace(0, 1, len(data[group_column].unique())))

        # 使用seaborn创建分组箱线图
        sns.boxplot(x=x_column, y=y_column, hue=group_column,
                    data=data, ax=ax, palette=colors)

        # 旋转X轴标签
        if rotation != 0:
            plt.setp(ax.get_xticklabels(), rotation=rotation, ha='right')

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title, fontsize=16, fontweight='bold')
    ax.legend(loc=legend_loc)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    return fig, ax