# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/99_Image_resize.ipynb (unless otherwise specified).

__all__ = ['display_image', 'glob_suffixes', 'resize_dataset']

# Cell

from pathlib import Path
from PIL import Image
from typing import Tuple, Optional, Dict, List, Union, Callable, Iterable

# Internal Cell

from fastcore.script import *

#import dolphins

import numpy as np
from shutil import rmtree

from .helpers import *
from tempfile import TemporaryDirectory

# Internal Cell


def _get_pallete(img_P: Image.Image) -> List[List[int]]:
    """Returns a list of colors from the input image"""
    palette = np.array(img_P.getpalette()).reshape((-1, 3)).tolist()
    palette = palette[: len(np.unique(palette, axis=0))]

    return palette

# Cell


def display_image(img: Union[Image.Image, Iterable[Image.Image]], x_size=240):
    """Displays an image or images with smaller dimension equal to x_size,
    and the other dimension is calculated"""

    if isinstance(img, Image.Image):
        y_size = int(img.size[1] / img.size[0] * x_size)
        display(img.resize((x_size, y_size)))
    else:
        for i in img:
            display_image(i, x_size=x_size)

# Internal Cell


def _separate_img(
    img: Image.Image,  # png image
    palette: List[List[int]],  # list of RGB colors
) -> List[Image.Image]:
    """Creates a list of images from the input image and color palette.
    Each image has one class and background"""

    img_array = np.array(img)
    img_separated = [
        Image.fromarray(img_array == i, mode="P") for i, color in enumerate(palette)
    ]

    black = [0, 0, 0]
    for color, im in zip(palette, img_separated):
        im.putpalette(black + color)

    return img_separated

# Internal Cell


def separate_masks(img: Image.Image) -> List[Image.Image]:
    """Converts RGB image to P. Creates a list of images from the input image.
    Each image has one class and background"""

    img_P = img.convert("P", palette=Image.ADAPTIVE, colors=len(img.getcolors()))

    palette = _get_pallete(img_P)

    return _separate_img(img_P, palette)

# Internal Cell


def resize_composite_multi(img: Image.Image, image_size: Tuple) -> Image.Image:
    """Resizes images and merges images into one"""
    separated_masks = separate_masks(img)
    separated_masks_new = []
    for im in separated_masks:
        separated_masks_new.append(im.resize(image_size))

    temp_img = Image.composite(
        separated_masks_new[1].convert("RGB"),
        separated_masks_new[0].convert("RGB"),
        separated_masks_new[1].convert("L"),
    )
    if len(separated_masks_new) <= 3:
        return temp_img
    else:
        for k in range(2, len(separated_masks_new) - 1):
            temp_img = Image.composite(
                separated_masks_new[k].convert("RGB"),
                temp_img.convert("RGB"),
                separated_masks_new[k]
                .convert("L")
                .point(lambda x: 0 if x == 0 else 255, "1"),
            )
        return temp_img

# Internal Cell


def calculate_image_size(img: Image.Image, smaller_size=1000) -> Tuple[int]:
    """Sets the smaller dimension of the image to the value of the variable smaller_size and reduces the larger dimension by the same amount"""
    x, y = img.size

    if min(x, y) == y:

        x = int(x / (y / smaller_size))
        y = smaller_size
    else:
        y = int(y / (x / smaller_size))
        x = smaller_size

    return (x, y)

# Internal Cell


def file_suffix_in(f: Path, suffixes: List[str]) -> bool:
    """Checks if the path contains a suffix"""
    return f.suffix[1:].lower() in [suffix.lower() for suffix in suffixes]

# Cell


def glob_suffixes(root_path: Path, suffixes: Union[List[str], str]) -> List[Path]:
    """Returns all suffixes located in the path and in the input variable suffixes"""
    if isinstance(suffixes, str):
        suffixes = [suffixes]
    return sorted([f for f in root_path.glob("**/*") if file_suffix_in(f, suffixes)])

# Internal Cell


def get_files_from_path(src_path: Path) -> (Path, Path, Path):
    """Returns three lists of files from the input directory"""
    src_images_path = Path(src_path / "JPEGImages/")
    src_class_path = Path(src_path / "SegmentationClass/")
    src_instance_path = Path(src_path / "SegmentationObject/")

    ensure_exists(src_images_path)
    ensure_exists(src_class_path)
    ensure_exists(src_instance_path)

    src_images = glob_suffixes(src_images_path, ("jpg", "jpeg"))
    src_lbl_classes = glob_suffixes(src_class_path, "png")
    src_lbl_instances = glob_suffixes(src_instance_path, "png")

    return src_images, src_lbl_classes, src_lbl_instances

# Internal Cell


def recreate_dst_dir(dst_path: Path) -> (Path, Path, Path):
    """Deletes and recreates the directory and returns 3 dst paths"""
    dst_images_path = Path(dst_path / "JPEGImages")
    dst_class_path = Path(dst_path / "SegmentationClass")
    dst_instance_path = Path(dst_path / "SegmentationObject")

    rmtree(dst_images_path, ignore_errors=True)
    rmtree(dst_class_path, ignore_errors=True)
    rmtree(dst_instance_path, ignore_errors=True)

    dst_images_path.mkdir(parents=True, exist_ok=True)
    dst_class_path.mkdir(parents=True, exist_ok=True)
    dst_instance_path.mkdir(parents=True, exist_ok=True)

    return dst_images_path, dst_class_path, dst_instance_path

# Internal Cell


def save_img_lbl_inst(
    src_img_file,
    src_class_file,
    src_instance_file,
    dst_images_path,
    dst_class_path,
    dst_instance_path,
    smaller_size=1000,
):
    """Resizes images from the src directories and saves them to the dst directories"""

    lbl = Image.open(src_class_file)
    image_size = calculate_image_size(lbl, smaller_size)
    lbl = resize_composite_multi(lbl, image_size)

    inst = Image.open(src_instance_file)
    inst = resize_composite_multi(inst, image_size)

    img = Image.open(src_img_file)
    img = img.resize(image_size)

    dst_img_file = dst_images_path / src_img_file.name
    dst_class_file = dst_class_path / src_class_file.name
    dst_instance_file = dst_instance_path / src_instance_file.name

    img.save(dst_img_file)
    lbl.save(dst_class_file)
    inst.save(dst_instance_file)

# Cell


@call_parse
def resize_dataset(
    src_path: Param("input directory containing images and annotations", Path),
    dst_path: Param("output directory, will be deleted if not empy", Path),
    dst_image_size: Param("min height or width", int) = 1000,
    verbose: Param("", bool) = True,
):

    """Resizes the entire dataset (images and labels) and saves it"""

    src_images, src_lbl_classes, src_lbl_instances = get_files_from_path(src_path)

    for src_img_path, src_class_path, src_instance_path in zip(
        src_images, src_lbl_classes, src_lbl_instances
    ):
        ensure_equal(src_img_path.stem, src_class_path.stem)
        ensure_equal(src_img_path.stem, src_instance_path.stem)

    dst_images_path, dst_class_path, dst_instance_path = recreate_dst_dir(dst_path)

    n = len(src_images)
    print(f"Resizing images and annotation masks in: {src_path.resolve()}")
    for i, (src_img_file, src_class_file, src_instance_file) in enumerate(
        zip(src_images, src_lbl_classes, src_lbl_instances)
    ):
        if verbose:
            print(f" - [{i+1}/{n}] {src_img_file.name}...")

        save_img_lbl_inst(
            src_img_file,
            src_class_file,
            src_instance_file,
            dst_images_path,
            dst_class_path,
            dst_instance_path,
        )