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


def violin_plot(data_list=None, labels=None, title="Violin Plot",
                colors=None, show_means=False, show_medians=True,
                show_extrema=True, vert=True):
    """
    绘制小提琴图

    参数:
    data_list: 数据列表
    labels: 数据标签
    title: 图表标题
    colors: 颜色列表
    show_means: 是否显示均值
    show_medians: 是否显示中位数
    show_extrema: 是否显示极值
    vert: 是否垂直显示
    """
    fig, ax = plt.subplots(figsize=(10, 6))

    # 生成默认数据
    if data_list is None:
        np.random.seed(42)
        data_list = [
            np.random.normal(0, 1, 100),
            np.random.normal(2, 1.5, 100),
            np.random.normal(-1, 0.5, 100)
        ]
        labels = ['Group A', 'Group B', 'Group C']

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

    # 绘制小提琴图
    parts = ax.violinplot(data_list,
                          showmeans=show_means,
                          showmedians=show_medians,
                          showextrema=show_extrema,
                          vert=vert)

    # 设置小提琴颜色
    for i, pc in enumerate(parts['bodies']):
        pc.set_facecolor(colors[i % len(colors)])
        pc.set_alpha(0.7)
        pc.set_edgecolor('black')

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

    if vert:
        ax.set_ylabel('Value', fontsize=12)
        ax.set_xlabel('Groups', fontsize=12)
        ax.set_xticks(np.arange(1, len(labels) + 1))
        ax.set_xticklabels(labels)
    else:
        ax.set_xlabel('Value', fontsize=12)
        ax.set_ylabel('Groups', fontsize=12)
        ax.set_yticks(np.arange(1, len(labels) + 1))
        ax.set_yticklabels(labels)

    ax.grid(True, linestyle='--', alpha=0.3, axis='y' if vert else 'x')

    plt.tight_layout()
    return fig, ax, parts