import matplotlib.pyplot as plt
import numpy as np
import os

def save_color_gradient(cmap_name='jet', width=512, height=50, filename=None):
    """
    保存颜色从 0 到 1 渐变的图像
    :param cmap_name: 色图名，如 'jet', 'viridis'
    :param width: 渐变图宽度（像素）
    :param height: 渐变图高度（像素）
    :param filename: 保存的文件名（可选）
    """
    cmap = plt.get_cmap(cmap_name)
    gradient = np.linspace(0, 1, width).reshape(1, -1)
    gradient = np.repeat(gradient, height, axis=0)


    plt.figure(figsize=(8, 1.5))
    plt.imshow(gradient, aspect='auto', cmap=cmap)
    plt.title(f"Color Gradient (colormap = '{cmap_name}')")
    plt.xlabel("Score: 0 → 1")
    plt.yticks([])
    plt.xticks([0, width//4, width//2, 3*width//4, width], [0.0, 0.25, 0.5, 0.75, 1.0])
    plt.tight_layout()

    # 默认文件名
    if filename is None:
        filename = f"color_gradient_{cmap_name}.png"

    # 保存图像
    plt.savefig(filename, dpi=150)
    plt.close()

    print(f"已保存: {os.path.abspath(filename)}")

# 调用：保存 jet 渐变图
save_color_gradient('jet')