# paint/stacked_bar_chart.py
import matplotlib.pyplot as plt
import numpy as np


def stacked_bar_chart(categories=None, data_matrix=None, labels=None,
                      title="Stacked Bar Chart", xlabel="Categories",
                      ylabel="Values", colors=None, edgecolor='black'):
    """
    绘制堆叠条形图

    参数:
    categories: 类别列表
    data_matrix: 数据矩阵 (n_categories, n_stacks)
    labels: 堆叠部分标签
    title: 图表标题
    xlabel: X轴标签
    ylabel: Y轴标签
    colors: 颜色列表
    edgecolor: 边框颜色
    """
    fig, ax = plt.subplots(figsize=(10, 6))

    # 生成默认数据
    if categories is None:
        categories = ['Q1', 'Q2', 'Q3', 'Q4']
    if data_matrix is None:
        np.random.seed(42)
        n_categories = len(categories)
        n_stacks = 4
        data_matrix = np.random.randint(10, 50, (n_categories, n_stacks))
        labels = ['Product A', 'Product B', 'Product C', 'Product D']

    n_categories = len(categories)
    n_stacks = data_matrix.shape[1]

    # 设置默认颜色
    if colors is None:
        colors = plt.cm.Paired(np.linspace(0, 1, n_stacks))

    # 计算累积高度用于堆叠
    cumulative = np.zeros(n_categories)

    # 绘制堆叠条形
    bars = []
    for i in range(n_stacks):
        bars.append(ax.bar(np.arange(n_categories), data_matrix[:, i],
                           bottom=cumulative,
                           color=colors[i % len(colors)],
                           edgecolor=edgecolor,
                           label=labels[i] if labels else f'Stack {i + 1}',
                           alpha=0.8))
        cumulative += data_matrix[:, i]

    # 设置图表属性
    ax.set_title(title, fontsize=14, fontweight='bold', pad=20)
    ax.set_xlabel(xlabel, fontsize=12)
    ax.set_ylabel(ylabel, fontsize=12)

    ax.set_xticks(np.arange(n_categories))
    ax.set_xticklabels(categories)

    ax.grid(True, linestyle='--', alpha=0.3, axis='y')

    # 添加数值标签
    for i in range(n_categories):
        total_height = 0
        for j in range(n_stacks):
            height = data_matrix[i, j]
            if height > 0:
                ax.text(i, total_height + height / 2,
                        str(height), ha='center', va='center',
                        fontsize=9, fontweight='bold', color='white')
            total_height += height

    # 添加图例
    ax.legend(fontsize=10, loc='upper left', bbox_to_anchor=(1.02, 1))

    plt.tight_layout()
    return fig, ax, bars