"""
热力图
"""

import matplotlib.pyplot as plt
import seaborn as sns
from vizpaint.utils import set_chinese_font


def create_heatmap(data, columns=None, title="热力图",
                   figsize=(12, 10), cmap='RdYlBu_r',
                   annot=True, fmt=".2f", cbar=True,
                   square=False):
    """
    创建热力图

    参数:
    ----------
    data : DataFrame 或 相关矩阵
        包含数据的DataFrame或相关矩阵
    columns : list, optional
        要计算相关的列名列表
    title : str, optional
        图表标题
    figsize : tuple, optional
        图表尺寸
    cmap : str, optional
        颜色映射
    annot : bool, optional
        是否在热力图上显示数值
    fmt : str, optional
        数值格式
    cbar : bool, optional
        是否显示颜色条
    square : bool, optional
        是否使热力图单元格为正方形

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

    if columns is not None:
        # 如果指定了列，则计算这些列的相关矩阵
        corr_matrix = data[columns].corr()
    else:
        # 否则直接使用data作为相关矩阵
        corr_matrix = data

    # 创建热力图
    sns.heatmap(corr_matrix, annot=annot, fmt=fmt, cmap=cmap,
                ax=ax, cbar=cbar, square=square,
                linewidths=0.5, linecolor='white')

    ax.set_title(title, fontsize=16, fontweight='bold')

    plt.tight_layout()
    return fig, ax