import os
import shutil
import httpx
import tqdm
from PIL import Image, ImageOps
import tempfile
from pathlib import Path
import random

VALIDATION_SPLIT = 0.1

query_yaml = """
want:
    - has: ["waste"]
    - nature: ["photo"]
"""

object_types = httpx.get("https://datasets.roundabout-host.com/api/object-types").json()

base_types = {
    "plastic household waste": set(),
    "glass household waste": set(),
    "metal household waste": set(),
    "paper and cardboard": set(),
    "organic household waste": set(),
    "household waste": set(),
    "waste": set(),
}

memoisation = {}

def get_objects(object_info):
    if object_info["id"] in memoisation:
        return memoisation[object_info["id"]]
    objects = set()
    for object in object_info["children"]:
        #print(f"Scanning {object}")
        info = httpx.get("https://datasets.roundabout-host.com/api/object/" + object).json()
        objects |= get_objects(info)
    objects.add(object_info["id"])
    memoisation[object_info["id"]] = objects
    return objects

for object_name in base_types:
    object_info = httpx.get("https://datasets.roundabout-host.com/api/object/" + object_name).json()
    base_types[object_name] = get_objects(object_info)
    print(object_name, base_types[object_name])

result = {"resources": True}   # dummy value to enter the loop
photos = []
offset = 0
limit = 192
output = Path("data/")
val_output = Path("val_data/")

print("Downloading photo metadata...")

while result["resources"]:
    result = httpx.post(f"https://datasets.roundabout-host.com/api/query-pictures?offset={offset}&limit={limit}", data={"query": query_yaml}).json()
    photos += result["resources"]
    offset += limit
    print(f"Received photos {offset-limit} to {offset}")

shutil.rmtree(output, ignore_errors=True)
os.makedirs(output)
shutil.rmtree(val_output, ignore_errors=True)
os.makedirs(val_output)

class_mapping = {
    "plastic household waste": 0,
    "glass household waste": 1,
    "metal household waste": 2,
    "paper and cardboard": 3,
    "organic household waste": 4,
    "household waste": 5,
    "waste": 6,
}

with tempfile.NamedTemporaryFile(delete_on_close=False, mode="wb") as temporary_file:
    for photo in tqdm.tqdm(photos):
        # Download the photo
        result = httpx.get(photo["download"], follow_redirects=True, timeout=15)
        temporary_file.write(result.content)
        temporary_file.seek(0)
        image = Image.open(temporary_file.name, formats=["JPEG"])
        ImageOps.exif_transpose(image, in_place=True)
        image.thumbnail((640, 640))
        if random.random() < VALIDATION_SPLIT:
            directory = val_output
        else:
            directory = output
        image.save(directory / (str(photo["id"]) + ".jpg"))
        # Download the annotations
        with open(directory / (str(photo["id"]) + ".txt"), "w") as annotation_file:
            for region in photo["regions"]:
                klass = -1
                for base_type in base_types:
                    if region["object"] in base_types[base_type]:
                        klass = class_mapping[base_type]
                        break
                if klass == -1:
                    continue

                if region["type"] == "bbox":
                    cx = region["shape"]["x"] + region["shape"]["w"] / 2
                    cy = region["shape"]["y"] + region["shape"]["h"] / 2
                    w = region["shape"]["w"]
                    h = region["shape"]["h"]
                elif region["type"] == "polygon":
                    x = [point["x"] for point in region["shape"]]
                    y = [point["y"] for point in region["shape"]]
                    cx = (min(x) + max(x)) / 2
                    cy = (min(y) + max(y)) / 2
                    w = max(x) - min(x)
                    h = max(y) - min(y)

                annotation_file.write(f"{klass} {cx} {cy} {w} {h}\n")