from io import BytesIO
import numpy as np
import os
from PIL import Image, ImageOps
import requests

# Local imports
from utils import Namespace


def download_from_urls(dictionary, root="."):
    for k, v in dictionary.items():
        directory = "/".join((root, k))
        if isinstance(v, list):
            os.makedirs(directory, exist_ok=True)
            print("Downloading images to %s" % directory)
            for i, item in ProgressBar(enumerate(v), prefix="Progress: "):
                # Download image
                try:
                    image = Image.open(BytesIO(requests.get(item["image_url"]).content))
                except:
                    msg = "\r  Unable to download image %i : %s"
                    print((msg % (i, item["image_url"])).ljust(70))
                    continue
                # Transform image
                if "transforms" in item.keys():
                    image = transform_image(image, item["transforms"])
                # Save image
                try:
                    image.save("%s/%s.png" % (directory, str(i).zfill(4)))
                except SystemError:
                    msg = "\r  Unable to save image %i : %s (there may be a problem with the image transforms)."
                    print((msg % (i, item["image_url"])).ljust(70))
        else:
            download_from_urls(v, root=directory)


def transform_image(image, transforms):
    for transform in transforms:
        if transform["name"].lower() == "crop":
            image = image.crop(**transform["kwargs"])
        elif transform["name"].lower() == "grayscale":
            image = ImageOps.grayscale(image)
        elif transform["name"].lower() == "mask":
            mask = np.array(Image.open(**transform["kwargs"]))
            image = np.array(image)
            mask[mask > 0] = 1
            if image.ndim == 3:
                mask = mask[..., None]
            image = Image.fromarray(image * mask)
    return image


class ProgressBar(object):

    def __init__(self, iterable, prefix="", length=50, fill="█"):
        self._iterable = list(iterable)
        self.prefix = prefix
        self.length = length
        self.fill = fill

    def __getitem__(self, i):
        if len(self._iterable) > 0:
            p = i / len(self._iterable)  # Progress as a decimal
            f = int(self.length * p)  # Fraction of the progress bar
            bar = self.fill * f + '-' * (self.length - f)  # Progress bar as a string
            line = "%s|%s| %.1f%%" % (self.prefix, bar, p * 100)  # Line to print
            print("\r%s" % line, end="\n" if i == len(self._iterable) else "\r")
        return self._iterable[i]


if __name__ == "__main__":
    import argparse
    import zipfile


    def unzip(path):
        print("Unzipping %s..." % path)
        with zipfile.ZipFile(path, "r") as f:
            f.extractall()
        print("Done")


    parser = argparse.ArgumentParser()
    parser.add_argument("-f", "--filepath", default="urls.yaml", type=str, help="Path to YAML file of URLs")
    parser.add_argument("-o", "--output", default="irships/augment", type=str, help="Output directory")
    parser.add_argument("-m", "--masks", default="masks", type=str, help="Directory of masks")
    args = parser.parse_args()

    if not os.path.isdir("irships/images") or not os.path.isdir("irships/labels"):
        unzip("irships.zip")  # Unzip IRShips if not already unzipped
    unzip(args.masks + ".zip")  # Unzip masks if not already unzipped

    download_from_urls(Namespace(args.filepath), root=args.output)