import pyvips
import glob
import random
import os
import math


def create_grid_composite(
    directory, k, spacing_inches, output_file="output.png", dpi=300
):
    # Calculate spacing in pixels
    spacing_pixels = int(spacing_inches * dpi)

    # Glob for PNG images
    png_files = glob.glob(os.path.join(directory, "*.png"))

    # Randomly select K^2 images
    selected_files = random.sample(png_files, k * k)

    # Create an empty list to hold the images
    images = [
        pyvips.Image.new_from_file(file, access="sequential") for file in selected_files
    ]

    # Calculate the size of the composite image
    widths, heights = zip(*[image.size for image in images])
    max_width = max(widths)
    max_height = max(heights)

    # Calculate total size of the grid including spacing
    total_width = k * max_width + (k + 1) * spacing_pixels
    total_height = k * max_height + (k + 1) * spacing_pixels

    # Create a blank image for the composite
    composite = pyvips.Image.black(total_width, total_height)

    # Place images into the composite
    for i, image in enumerate(images):
        row = i // k
        col = i % k
        x = col * (max_width + spacing_pixels) + spacing_pixels
        y = row * (max_height + spacing_pixels) + spacing_pixels
        composite = composite.insert(image, x, y)

    # Save the composite image
    composite.write_to_file(output_file)


if __name__ == "__main__":
    # Example usage
    directory = "/teamspace/studios/this_studio/out_sortcolors/umap/"  # Change to your directory path
    k = 3  # For a 3x3 grid
    spacing_inches = 0.5  # Half an inch between images
    create_grid_composite(directory, k, spacing_inches)