import sys
from pathlib import Path
from skimage.io import imread, imsave
from skimage.transform import resize
from skimage.color import rgba2rgb
from argparse import ArgumentParser
import numpy as np

IMG_EXTENSIONS = set(
    [".jpg", ".JPG", ".jpeg", ".JPEG", ".png", ".PNG", ".ppm", ".PPM", ".bmp", ".BMP"]
)


def is_image_file(filename):
    """Check that a file's name points to a known image format
    """
    if isinstance(filename, Path):
        return filename.suffix in IMG_EXTENSIONS

    return Path(filename).suffix in IMG_EXTENSIONS


def find_images(path, recursive=False):
    """
    Get a list of all images contained in a directory:

    - path.glob("*") if not recursive
    - path.glob("**/*") if recursive
    """
    p = Path(path)
    assert p.exists()
    assert p.is_dir()
    pattern = "*"
    if recursive:
        pattern += "*/*"

    return [i for i in p.glob(pattern) if i.is_file() and is_image_file(i)]


def uint8(array):
    return array.astype(np.uint8)


def crop_and_resize(image_path, label_path):
    """
    Resizes an image so that it keeps the aspect ratio and the smallest dimensions
    is 640, then crops this resized image in its center so that the output is 640x640
    without aspect ratio distortion

    Args:
        image_path (Path or str): Path to an image
        label_path (Path or str): Path to the image's associated label

    Returns:
        tuple((np.ndarray, np.ndarray)): (new image, new label)
    """
    dolab = label_path is not None

    img = imread(image_path)
    if dolab:
        lab = imread(label_path)

    if img.shape[-1] == 4:
        img = uint8(rgba2rgb(img) * 255)

    if dolab and img.shape != lab.shape:
        print("\nWARNING: shape mismatch. Entering breakpoint to investigate:")
        breakpoint()

    # resize keeping aspect ratio: smallest dim is 640
    h, w = img.shape[:2]
    if h < w:
        size = (640, int(640 * w / h))
    else:
        size = (int(640 * h / w), 640)

    r_img = resize(img, size, preserve_range=True, anti_aliasing=True)
    r_img = uint8(r_img)

    if dolab:
        # nearest neighbor for labels
        r_lab = resize(lab, size, preserve_range=True, anti_aliasing=False, order=0)
        r_lab = uint8(r_lab)

    # crop in the center
    H, W = r_img.shape[:2]

    top = (H - 640) // 2
    left = (W - 640) // 2

    rc_img = r_img[top : top + 640, left : left + 640, :]
    if dolab:
        rc_lab = r_lab[top : top + 640, left : left + 640, :]
    else:
        rc_lab = None

    return rc_img, rc_lab


def label(img, label, alpha=0.4):
    return uint8(alpha * label + (1 - alpha) * img)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "-i", "--input_dir", type=str, help="Directory to recursively read images from"
    )
    parser.add_argument(
        "-o",
        "--output_dir",
        type=str,
        help="Where to writ the result of the script,"
        + " keeping the input dir's structure",
    )
    parser.add_argument(
        "--no_labels",
        action="store_true",
        help="Only process images, don't look for labels",
    )
    parser.add_argument(
        "--store_labeled",
        action="store_true",
        help="Store a superposition of the label and the image in out/labeled/",
    )
    args = parser.parse_args()

    dolab = not args.no_labels
    dolabeled = args.store_labeled

    input_base = Path(args.input_dir).expanduser().resolve()
    output_base = Path(args.output_dir).expanduser().resolve()

    input_images = input_base / "imgs"
    output_images = output_base / "imgs"

    if dolab:
        input_labels = input_base / "labels"
        output_labels = output_base / "labels"
        if dolabeled:
            output_labeled = output_base / "labeled"

    print("Input images:", str(input_images))
    print("Output images:", str(output_images))
    if dolab:
        print("Input labels:", str(input_labels))
        print("Output labels:", str(output_labels))
        if dolabeled:
            print("Output labeled:", str(output_labeled))
    else:
        print("NO LABEL PROCESSING (args.no_labels is specified)")
    print()

    assert input_images.exists()
    if dolab:
        assert input_labels.exists()

    if output_base.exists():
        if (
            "n"
            in input(
                "WARNING: output dir already exists."
                + " Overwrite its content? (y/n, default: y)"
            ).lower()
        ):
            sys.exit()

    output_images.mkdir(parents=True, exist_ok=True)
    if dolab:
        output_labels.mkdir(parents=True, exist_ok=True)
        if dolabeled:
            output_labeled.mkdir(parents=True, exist_ok=True)

    images_paths = list(
        map(Path, sorted((map(str, find_images(input_images, recursive=True)))))
    )
    if dolab:
        labels_paths = list(
            map(Path, sorted((map(str, find_images(input_labels, recursive=True)))))
        )
    else:
        labels_paths = [None] * len(images_paths)

    for i, (image_path, label_path) in enumerate(zip(images_paths, labels_paths)):
        print(
            f"Processing {i + 1 :3} / {len(images_paths)} : {image_path.name}",
            end="\r",
            flush=True,
        )
        processed_image, processed_label = crop_and_resize(image_path, label_path)
        imsave(output_images / f"{image_path.stem}.png", processed_image)
        if dolab:
            imsave(output_labels / f"{label_path.stem}.png", processed_label)
            if dolabeled:
                labeled = label(processed_image, processed_label)
                imsave(output_labeled / f"{image_path.stem}.png", labeled)

    print("\nDone.")