# Knowledgedump.org - random_shape_generation
# Creates 1000 images of shapes, consisting of triangles, ellipses and rectangles.
# The shapes are white on a black 800x800 pixel background by default.
# Each shape has minimum and maximum size requirements, so that they can't be too small/large.

# Required packages: numpy, matplotlib, scikit-image.

import numpy as np
import matplotlib.pyplot as pypl
import os
import random
from skimage.draw import polygon, ellipse
import zipfile


# Function to generate a random "well-formed" triangle (not too small/large/close to a line)
def generate_triangle(image_size=(800, 800)):
    # Maximum and minimum sizes are 20% and 5% of the smallest img dimension
    max_size = round(min(image_size) * 0.2)
    min_size = round(min(image_size) * 0.05)
    
    # Randomly choose first point
    x1, y1 = random.randint(0, image_size[0]), random.randint(0, image_size[1])
    
    # Randomly choose second point within distance of [min_size, max_size]
    x2, y2 = random.randint(max(0, x1 - max_size), min(image_size[0], x1 + max_size)), random.randint(max(0, y1 - max_size), min(image_size[1], y1 + max_size))
    
    # To ensure that the triangle is not misshapen, the third point can't be too close to the line formed
    # between points 1 and 2. This can for instance be implemented, by calculating the orthogonal vector
    # to the vector between pts 1 and 2, normalizing it and then applying the size requirement for the "height"
    # of our triangle. The third point is then chosen as the vector we thus constructed, added to some point on the line
    # between points 1 and 2 (here, we simply pick the midpoint between 1 and 2).

    # Calculate the orthogonal vector to the vector between points 1 and 2, given by np.array([x2 - x1, y2 - y1])
    v_orth = np.array([y1 - y2, x2 - x1])
    
    # Normalize the orthogonal
    v_orth_norm = v_orth / np.linalg.norm(v_orth)
    
    # Random distance for the third point from the line between points 1 and 2 (within parameters)
    height = random.randint(min_size, max_size)
    
    # Calculate midpoint and place third point.
    midpoint = np.array([x1 + (x2 - x1) / 2, y1 + (y2 - y1) / 2])
    x3, y3 = midpoint + height * v_orth_norm
    
    # Ensure the third point is within the image bounds. This could still result in "misshapen" triangles,
    # when point 1 and 2 are close to the image border and the orthogonal vector points "outside",
    # but it's unlikely that there will be many of them, so we don't bother adding more complexity to this
    # and simply ensure that point 3 is inside the image borders.
    x3, y3 = int(np.clip(x3, 0, image_size[0])), int(np.clip(y3, 0, image_size[1]))
    
    # Initialize a black image (array of zeros)
    image = np.zeros(image_size, dtype=np.uint8)
    
    # Create the triangle using the polygon function from skimage.draw (scikit-image).
    # It returns the pixel coordinates for the triangle, so we can set those to white (255) on our black image.
    # (see pypl.imsave colormap settings below)
    rr, cc = polygon([y1, y2, y3], [x1, x2, x3], shape=image_size)
    image[rr, cc] = 255
    
    return image

# Function to generate a random ellipse with a size cap and minimum size
def generate_ellipse(image_size=(800, 800)):
    # Maximum and minimum diameters are set to 20% and 5% of the smallest img dimension
    max_radius = round(min(image_size) * 0.1)
    min_radius = round(min(image_size) * 0.025)
    
    # Set random horizontal and vertical radius of ellipse
    rx = random.randint(min_radius, max_radius)
    ry = random.randint(min_radius, max_radius)
    
    # Random center
    cx = random.randint(rx, image_size[0] - rx)
    cy = random.randint(ry, image_size[1] - ry)
    
    # Initialize a black image
    image = np.zeros(image_size, dtype=np.uint8)
    
    # Create the ellipse using the ellipse function from skimage.draw, setting the values to 255 for white.
    rr, cc = ellipse(cx, cy, rx, ry)
    image[rr, cc] = 255
    
    return image

# Function to generate a random rectangle within size limits.
def generate_rectangle(image_size=(800, 800)):
    # Maximum and minimum sizes are 20% and 5% of the smallest img dimension
    max_size = round(min(image_size) * 0.2)
    min_size = round(min(image_size) * 0.05)
    
    # Random width and height within the size range
    rect_width = random.randint(min_size, max_size)
    rect_height = random.randint(min_size, max_size)

    # Random position for the top-left corner of the rectangle
    x1 = random.randint(0, image_size[0] - rect_width)
    y1 = random.randint(0, image_size[1] - rect_height)

    # Initialize a black image
    image = np.zeros(image_size, dtype=np.uint8)

    # Draw the rectangle (fill the area) in white colour (255).
    image[y1:y1+rect_height, x1:x1+rect_width] = 255
    
    return image



if __name__ == "__main__":
    # Create folder at the location of the python script to save images.
    current_dir = os.path.dirname(os.path.abspath(__file__))
    output_dir = os.path.join(current_dir, "random_segments")
    os.makedirs(output_dir, exist_ok=True)


    # Function to generate and save randomly shaped segments (triangles, ellipses, rectangles).
    def generate_random_segments(image_size=(800, 800), num_segments=1000):
        image_files = []
        for i in range(num_segments):
            shape_type = random.choice(["triangle", "ellipse", "rectangle"])
        
            if shape_type == "triangle":
                segment = generate_triangle(image_size)
            elif shape_type == "ellipse":
                segment = generate_ellipse(image_size)
            else:
                segment = generate_rectangle(image_size)
        
            # Save the segment as a PNG image, using matplotlib.imsave function with "gray" colormap.
            # The grayscale colormap goes from 0 for black to 255 for white.
            image_filename = os.path.join(output_dir, f"segment_{i+1}.png")
            pypl.imsave(image_filename, segment, cmap="gray", format="png")
            image_files.append(image_filename)
    
        # Return the list of saved image filenames
        return image_files


    # Generate num_segment (by default 1000) random segments and save them as images.
    image_files = generate_random_segments()

    # Compress the images into a zip file.
    zip_filename = os.path.join(current_dir, "random_segments.zip")
    # Using "with" keyword for cleaner opening the file and closing it after all files are written.
    # Call class constructor to create "Zipfile" object and add images, using its ".write" method.
    with zipfile.ZipFile(zip_filename, "w") as zipfi:
        for file in image_files:
            zipfi.write(file, os.path.basename(file))

    print(f"All images have been successfully created and compressed into {zip_filename}.")