"""
组合图（多种图表组合）
"""

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

def create_combination_chart(data, x_column, y_columns, title="组合图表",
                             figsize=(15, 12), colors=None):
    """
    创建组合图表（包含多种图表类型）

    参数:
    ----------
    data : DataFrame
        包含数据的DataFrame
    x_column : str
        X轴列名
    y_columns : list
        Y轴列名列表
    title : str, optional
        图表标题
    figsize : tuple, optional
        图表尺寸
    colors : list, optional
        颜色列表

    返回:
    ----------
    fig : 图表对象
    """
    set_chinese_font()

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

    fig = plt.figure(figsize=figsize)

    # 创建2x2的网格
    gs = GridSpec(2, 2, figure=fig)

    # 1. 多线图
    ax1 = fig.add_subplot(gs[0, 0])
    for i, col in enumerate(y_columns):
        ax1.plot(data[x_column], data[col], 'o-',
                 label=col, color=colors[i], linewidth=2)
    ax1.set_title('多线图')
    ax1.set_xlabel(x_column)
    ax1.set_ylabel('数值')
    ax1.legend(loc='upper left')
    ax1.grid(True, alpha=0.3)
    plt.setp(ax1.get_xticklabels(), rotation=45, ha='right')

    # 2. 分组柱状图
    ax2 = fig.add_subplot(gs[0, 1])
    x = range(len(data[x_column]))
    bar_width = 0.8 / len(y_columns)
    for i, col in enumerate(y_columns):
        offset = (i - len(y_columns) / 2 + 0.5) * bar_width
        ax2.bar([pos + offset for pos in x], data[col],
                width=bar_width, label=col, color=colors[i], alpha=0.8)
    ax2.set_title('分组柱状图')
    ax2.set_xlabel(x_column)
    ax2.set_ylabel('数值')
    ax2.set_xticks(x)
    ax2.set_xticklabels(data[x_column], rotation=45)
    ax2.legend(loc='upper right')
    ax2.grid(axis='y', alpha=0.3)

    # 3. 堆叠面积图
    ax3 = fig.add_subplot(gs[1, 0])
    ax3.stackplot(data[x_column], [data[col] for col in y_columns],
                  labels=y_columns, colors=colors, alpha=0.7)
    ax3.set_title('堆叠面积图')
    ax3.set_xlabel(x_column)
    ax3.set_ylabel('数值')
    ax3.legend(loc='upper left')
    ax3.grid(True, alpha=0.3)
    plt.setp(ax3.get_xticklabels(), rotation=45, ha='right')

    # 4. 散点图矩阵（简化版）
    ax4 = fig.add_subplot(gs[1, 1])
    # 只选择前两个系列绘制散点图
    if len(y_columns) >= 2:
        scatter = ax4.scatter(data[y_columns[0]], data[y_columns[1]],
                              c=data[y_columns[0]] if len(y_columns) > 2 else colors[0],
                              s=100, alpha=0.7, cmap='viridis')
        ax4.set_xlabel(y_columns[0])
        ax4.set_ylabel(y_columns[1])
        if len(y_columns) > 2:
            plt.colorbar(scatter, ax=ax4, label=y_columns[0])
    ax4.set_title('散点图')
    ax4.grid(True, alpha=0.3)

    fig.suptitle(title, fontsize=18, fontweight='bold', y=0.98)
    plt.tight_layout()

    return fig