from PIL import Image
import os
import math

def split_and_merge_image(image_path, chunk_width, chunk_height, output_directory):
    # Open the input image
    img = Image.open(image_path)

    # Get the dimensions of the input image
    img_width, img_height = img.size

    # Calculate the number of chunks in rows and columns
    num_chunks_x = img_width // chunk_width
    num_chunks_y = img_height // chunk_height

    # Calculate the total number of chunks
    total_chunks = num_chunks_x * num_chunks_y

    # Calculate the number of columns for the grid based on total chunks
    columns = int(math.ceil(math.sqrt(total_chunks)))

    # Calculate the number of rows based on the calculated columns
    rows = int(math.ceil(total_chunks / columns))

    # Calculate the size of the output squared canvas
    canvas_width = columns * chunk_width
    canvas_height = rows * chunk_height
    canvas = Image.new('RGB', (canvas_width, canvas_height))

    # Create the output directory if it doesn't exist
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    # Split the image into chunks and arrange them on the canvas
    for i in range(total_chunks):
        chunk_x = i % num_chunks_x
        chunk_y = i // num_chunks_x

        left = chunk_x * chunk_width
        upper = chunk_y * chunk_height
        right = left + chunk_width
        lower = upper + chunk_height

        chunk = img.crop((left, upper, right, lower))

        # Calculate the paste position on the canvas
        paste_x = (i % columns) * chunk_width
        paste_y = (i // columns) * chunk_height

        # Paste the chunk onto the canvas
        canvas.paste(chunk, (paste_x, paste_y))

    # Save the squared canvas image
    output_path = os.path.join(output_directory, "squared_canvas.png")
    canvas.save(output_path)

# Set the input image path, chunk width, chunk height, and output directory
input_image_path = "long.png"
chunk_width = 2098	  # Set the desired chunk width
chunk_height = 1400  # Set the desired chunk height
output_directory = "output_squared"

# Call the function to split and merge the image into a squared canvas
split_and_merge_image(input_image_path, chunk_width, chunk_height, output_directory)