from typing import Sequence
import matplotlib.pyplot as plt


def show_ims(
    ims: Sequence[Image.Image],
    lbls: Sequence[str] | None = None,
    row: int = 2,
    col: int = 5,
    cmap: str = "gray",
):
    n_data = row * col
    
    fig, ax = plt.subplots(nrows=row, ncols=col, figsize=(8,3))
    for i, im in enumerate(ims[:n_data]):
        r = i // col
        c = i % col
        if row == 1:
            if lbls is not None:
                ax[c].set_title(lbls[i], fontsize=8)
            ax[c].axes.xaxis.set_visible(False)
            ax[c].axes.yaxis.set_visible(False)
            ax[c].imshow(im, cmap=cmap)
        else:
            if lbls is not None:
                ax[r, c].set_title(lbls[i], fontsize=8)
            ax[r, c].axes.xaxis.set_visible(False)
            ax[r, c].axes.yaxis.set_visible(False)
            ax[r, c].imshow(im, cmap=cmap)

xx = np.linspace(0, np.pi/2, 3)
yy = np.linspace(0, np.pi/2, 10)
X, Y = np.meshgrid(xx, yy)
print(f"{X.shape=} {Y.shape=}")
pprint.pprint(X)
pprint.pprint(Y)

Z = np.sin(X) * np.cos(Y)

xy_pairs = zip(X.flatten(), Y.flatten())
Z_values = []
for x, y in xy_pairs:
    z = np.sin(x) * np.cos(y)
    Z_values.append(z)

Z_new = np.array(Z_values).reshape(X.shape)
print(np.allclose(Z, Z_new))

if True:
    if False:
        plt.pcolormesh(X, Y, Z, cmap='coolwarm')
        plt.colorbar()
        plt.xlabel('X-axis')
        plt.ylabel('Y-axis')
        plt.title('Heatmap of f(x,y)')
        plt.show()
    else:
        fig, ax = plt.subplots()
        c = ax.pcolormesh(X, Y, Z, cmap='coolwarm')
        fig.colorbar(c, ax=ax)

        # ラベルとタイトルの設定
        ax.set_xlabel('X-axis')
        ax.set_ylabel('Y-axis')
        ax.set_title('Heatmap of f(x,y)')

        # Figureをバッファに保存
        buf = BytesIO()
        fig.savefig(buf, format='png')
        plt.close(fig)  # 不要になったFigureを閉じる

        # バッファからPIL.Imageを作成
        buf.seek(0)
        im = Image.open(buf)
        display(im)

fig, ax = plt.subplots()
print(fig.bbox)

fig, ax = plt.subplots(figsize=(6.4, 4.8))