"""
堆叠柱状图
"""

import numpy as np
import matplotlib.pyplot as plt
from vizpaint.utils import set_chinese_font


def create_stacked_bar_chart(data, x_column, y_columns, title="堆叠柱状图",
                             xlabel="X轴", ylabel="Y轴", figsize=(12, 8),
                             colors=None, legend_loc='upper right', show_values=False,
                             rotation=45):
    """
    创建堆叠柱状图

    参数:
    ----------
    data : DataFrame
        包含数据的DataFrame
    x_column : str
        X轴列名
    y_columns : list
        要堆叠显示的Y轴列名列表
    title : str, optional
        图表标题
    xlabel : str, optional
        X轴标签
    ylabel : str, optional
        Y轴标签
    figsize : tuple, optional
        图表尺寸
    colors : list, optional
        颜色列表
    legend_loc : str, optional
        图例位置
    show_values : bool, optional
        是否在柱状图上显示数值
    rotation : int, optional
        X轴标签旋转角度

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

    x = np.arange(len(data[x_column]))
    bottom = np.zeros(len(data[x_column]))

    if colors is None:
        colors = plt.cm.Set3(np.linspace(0, 1, len(y_columns)))

    for i, col in enumerate(y_columns):
        bars = ax.bar(x, data[col], bottom=bottom, label=col,
                      color=colors[i], edgecolor='black', alpha=0.8)
        bottom += data[col].values

        if show_values:
            for j, bar in enumerate(bars):
                height = bar.get_height()
                if height > 0:
                    ax.text(bar.get_x() + bar.get_width() / 2.,
                            bottom[j] - height / 2,
                            f'{height:.1f}', ha='center', va='center',
                            fontsize=9, color='white' if height > bottom[j] / 2 else 'black')

    ax.set_xticks(x)
    ax.set_xticklabels(data[x_column], rotation=rotation)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title, fontsize=16, fontweight='bold')
    ax.legend(loc=legend_loc)
    ax.grid(axis='y', alpha=0.3)

    plt.tight_layout()
    return fig, ax