import pyvips
import glob
import random
import os
import math

import numpy as np


def bg_fill(width: int, height: int, channels: int = 3, fill: int = 255):
    a = np.ones(shape=(height, width, channels)) * fill
    return pyvips.Image.new_from_array(a)


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])
    widths, heights = [1800], [1800]
    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, bands=4)
    composite = bg_fill(total_width, total_height, channels=1, fill=255)

    # 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.flatten(background=[255, 255, 255]), x, y)

    # Save the composite image
    composite.write_to_file(output_file)
    x = pyvips.Image.thumbnail(output_file, 1080 * 4)
    x.write_to_file(output_file.replace(".png", "_pre.png"))
    print(output_file)


if __name__ == "__main__":
    # Example usage
    experiment = "umap"
    # experiment = "unsupervised-anchors"
    directory = f"/teamspace/studios/this_studio/out_sortcolors/{experiment}/"  # Change to your directory path
    # directory = f"/teamspace/studios/this_studio/colors/colors-refactor-{experiment}/"  # Change to your directory path
    k = 10
    spacing_inches = 2
    sp_str = str(spacing_inches).replace(".", "-")
    # savefile = f"/teamspace/studios/this_studio/composites/{experiment}-{sp_str}.png"
    savefile = f"/teamspace/studios/this_studio/composites/{experiment}-{sp_str}.png"
    create_grid_composite(directory, k, spacing_inches, output_file=savefile)