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


def heatmap(data=None, row_labels=None, col_labels=None,
            title="Heatmap", cmap='viridis', annotate=True,
            fmt='.2f', cbar=True, square=False):
    """
    绘制热力图

    参数:
    data: 二维数据数组
    row_labels: 行标签
    col_labels: 列标签
    title: 图表标题
    cmap: 颜色映射
    annotate: 是否显示数值
    fmt: 数值格式
    cbar: 是否显示颜色条
    square: 是否保持方形
    """
    fig, ax = plt.subplots(figsize=(10, 8))

    # 生成默认数据
    if data is None:
        np.random.seed(42)
        data = np.random.rand(8, 10)
        row_labels = [f'Row {i + 1}' for i in range(8)]
        col_labels = [f'Col {i + 1}' for i in range(10)]

    # 绘制热力图
    im = ax.imshow(data, cmap=cmap, aspect='auto' if not square else 'equal')

    # 添加数值标签
    if annotate:
        for i in range(data.shape[0]):
            for j in range(data.shape[1]):
                text = ax.text(j, i, format(data[i, j], fmt),
                               ha="center", va="center",
                               color="white" if data[i, j] > 0.6 else "black",
                               fontsize=8)

    # 设置刻度标签
    if row_labels is not None:
        ax.set_yticks(np.arange(data.shape[0]))
        ax.set_yticklabels(row_labels)
    if col_labels is not None:
        ax.set_xticks(np.arange(data.shape[1]))
        ax.set_xticklabels(col_labels, rotation=45, ha='right')

    # 设置标题
    ax.set_title(title, fontsize=14, fontweight='bold', pad=20)

    # 添加颜色条
    if cbar:
        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

    plt.tight_layout()
    return fig, ax, im