# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['get_img_files', 'resize_img', 'stack_imgs', 'img_save_path']

# %% ../nbs/00_core.ipynb 3
# Import necessary modules from the standard library
from pathlib import Path  # For working with file paths
import hashlib
from glob import glob
import numpy as np  # For working with arrays
from PIL import Image  # For working with images

# %% ../nbs/00_core.ipynb 5
def get_img_files(img_dir:Path, # The directory to search for image files
                    img_fmts=['jpg','jpeg','png'] # The list of image formats to search for
                   ):
    """
    Get all the image files in the given directory.
    
    Returns:
    img_paths (list): A list of pathlib.Path objects representing the image files
    """
    img_paths = []

    # Use the glob module to search for image files with specified formats
    for fmt in img_fmts:
        img_paths.extend(glob(f'{img_dir}/*.{fmt}'))
    # Convert the file paths to pathlib.Path objects
    img_paths = [Path(path) for path in img_paths]
    
    return img_paths

# %% ../nbs/00_core.ipynb 11
def resize_img(img:'Image', # The image to be resized
               target_sz:int=512, # The target size of the image
               divisor:int=32 # The divisor value to crop the image
              ):
    """
    Resize the image to the target size, keeping aspect ratio 
    and crop the image if the size is not divisible by divisor.
    
    Returns:
    img (PIL.Image): The resized and possibly cropped image
    """
    # Find the minimum and maximum dimension of the image
    min_dim = np.argmin(img.size)
    max_dim = np.argmax(img.size)
    
    # Calculate the ratio of the minimum dimension to the target size
    ratio = min(img.size)/target_sz
    
    new_sz = []
    
    # Insert the target size to the minimum dimension of the image
    new_sz.insert(min_dim, target_sz)
    
    # Insert the maximum dimension of the image calculated based on the ratio
    new_sz.insert(max_dim, int(max(img.size)/ratio))
    
    # Resize the image to the new size
    img = img.resize(new_sz)
    
    # Crop the image if divisor is greater than 0
    if divisor > 0:
        src_w, src_h = img.size
        width = src_w if src_w%divisor==0 else src_w - src_w%divisor
        height = src_h if src_h%divisor==0 else src_h - src_h%divisor
        img = img.crop(box=(0, 0, width, height))
    
    return img

# %% ../nbs/00_core.ipynb 17
def stack_imgs(imgs:['Image'] # A list of PIL.Image objects.
              ):
    """
    Stacks a list of images horizontally or vertically, depending on which dimension is larger.
    
    Returns:
    stack (PIL.Image): A single image containing all the input images stacked horizontally or vertically.
    """
    # Get the size of the first image in the list
    w,h = imgs[0].size
    
    # If the width is larger than the height, stack the images vertically
    if w >= h:
        # Create a new RGB image with the appropriate size
        stack = Image.new('RGB', size=(w, len(imgs)*h))
        
        # Iterate through the list of images and paste each image into the stack
        for i, img in enumerate(imgs): 
            stack.paste(img, box=(0, i*h, w, (i+1)*h))
    
    # If the height is larger than the width, stack the images horizontally
    else:
        # Create a new RGB image with the appropriate size
        stack = Image.new('RGB', size=(len(imgs)*w, h))
        
        # Iterate through the list of images and paste each image into the stack
        for i, img in enumerate(imgs): 
            stack.paste(img, box=(i*w, 0, (i+1)*w, h))
    
    # Return the resulting stacked image
    return stack

# %% ../nbs/00_core.ipynb 21
def img_save_path(image, # The image to be saved
                  save_dir, # The directory where the image will be saved
                  suffix='', # A suffix to add to the file name
                  hashlen=8, # The length of the image hash to use in the file name
                  fmt='PNG' # The format of the image
                 ):
    """
    Generate the file path to save the image.
    
    Returns:
    path (pathlib.Path): The file path to save the image
    """
    # Add a dash to the front of the suffix to the file name if it is provided
    # and hashlen is greater than 0
    suffix = f'-{suffix}' if len(suffix) > 0 and hashlen > 0 else suffix
    
    # Generate the SHA-256 hash of the image
    img_hash = hashlib.sha256(np.array(image)).hexdigest()[:hashlen]
    
    # Generate the file path by combining the save directory, image hash, suffix, and format
    return Path(f"{save_dir}/{img_hash}{suffix}.{fmt.lower()}")
