import os
import math
import argparse
import matplotlib.pyplot as plt
from PIL import Image

def on_key(event):
    if event.key in ['q', 'Q', 'escape']:
        plt.close()

def tile_images(folder, max_columns=10, image_size=(64, 64)):
    # PNGファイル一覧を取得
    image_paths = sorted([
        os.path.join(folder, f)
        for f in os.listdir(folder)
        if f.lower().endswith(".png")
    ])

    if not image_paths:
        print("No PNG images found.")
        return

    num_images = len(image_paths)
    columns = min(max_columns, num_images)
    rows = math.ceil(num_images / columns)

    # 普通の解像度
    fig, axes = plt.subplots(rows, columns, figsize=(columns, rows))
    fig.canvas.mpl_connect("key_press_event", on_key)
    plt.subplots_adjust(wspace=0.05, hspace=0.05)

    axes = axes.flatten() if num_images > 1 else [axes]

    for ax, path in zip(axes, image_paths):
        img = Image.open(path).convert("RGB")
        ax.imshow(img)
        ax.set_title(os.path.basename(path), fontsize=3)
        ax.axis("off")

    for i in range(len(image_paths), len(axes)):
        axes[i].axis("off")

    print("Press Q or ESC to close.")
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Tile display of PNG images in a folder.")
    parser.add_argument("folder", type=str, help="Folder containing PNG files")
    parser.add_argument("--max_columns", type=int, default=10, help="Maximum number of columns in the grid")
    args = parser.parse_args()

    tile_images(args.folder, max_columns=args.max_columns)