import json
import os
from pprint import pprint

from itertools import product
from typing import List

import cv2
import numpy as np
import pandas as pd
from PIL import Image, ImageDraw
from detectron2.structures import BoxMode
from shapely.geometry import Polygon, box, MultiPolygon
import geopandas as gpd

Image.MAX_IMAGE_PIXELS = None

# with open("../data/GeoJson/Flaechenbelaege.json", "r") as file:
#     flaechenbelaege = json.load(file)
#     print(flaechenbelaege)

flaechenbelaege = gpd.read_file("../data/GeoJson/Flaechenbelaege.json")

reference_points = pd.read_csv("../data/Referenzpunkte.csv")

dir_in = "../data/images/"
main_image_filename = "WestendDOP2.tif"

name, ext = os.path.splitext(main_image_filename)
img = Image.open(os.path.join(dir_in, main_image_filename))
w, h = img.size
corners = reference_points.head(4)
left = min(corners["X"])
right = max(corners["X"])
bottom = min(corners["Y"])
top = max(corners["Y"])
x_scale = (right - left) / w
y_scale = (bottom - top) / h

print(left, top)
print(right, bottom)

print("Creating canvas...")
draw = ImageDraw.Draw(img)
d = 1000
grid = product(range(0, h - h % d, d), range(0, w - w % d, d))

images = []
category_translations = {
    'Asphalt': "Festweg",
    'Bepflanzte_Flaechen': "Baumbestand",
    'Beton': "Festweg",
    'Noppenpflaster': "Pflaster",
    'Pflaster': "Pflaster",
    'Platten': "Pflaster",
    'Rasen': "Wiese",
    'Rasengittersteine': "Pflaster",
    'Rippenpflaster': "Pflaster",
    'Sand': "Festweg",
    'unbefestigt': "Festweg",
}
categories = [
    "Baumbestand",
    "Festweg",
    "Pflaster",
    "Wiese",
    "Wasser",
    "Gullydeckel",
]

for i, j in grid:
    # if i > 14000:
    #     break

    view_box = (
        j,
        i,
        j + d,
        i + d,
    )
    annotations = []
    filename = f"data/images/westend/belag/{i}_{j}.tif"
    image = dict(
        image_id=len(images),
        width=d,
        height=d,
        file_name=f"data/images/cropped/WestendDOP2_{i}_{j}.tif",
        annotations=annotations,
    )

    # print(flaechenbelaege.head())
    # print(flaechenbelaege.columns)

    # loop over every shape
    for _, row in flaechenbelaege.iterrows():
        kind = row["Art"]
        category_name = category_translations.get(kind, "")
        category_id = categories.index(category_name)
        first_obj: Polygon = row["geometry"]  # ["coordinates"][0]
        # pprint(first_obj)
        # print(first_obj.bounds)
        # polygon = Polygon(first_obj)
        # print(polygon)
        # print(polygon.bounds)
        # print(list(first_obj.exterior.coords))

        scaled = list(map(lambda xy: ((xy[0] - left) / x_scale, (xy[1] - top) / y_scale),
                          first_obj.exterior.coords))

        # pprint(scaled[:10])
        belag = Polygon(scaled)
        # print(belag.bounds)
        intersection = belag.intersection(box(*view_box))
        if not intersection.bounds:
            # print(f"No shape found in {i, j}")
            continue
        # print()
        # print(i, j)
        # print(intersection.bounds)
        # exit()

        # print("Drawing polygon...")
        if isinstance(intersection, MultiPolygon):
            intersects: List[Polygon] = intersection.geoms
            # print(f"error at {a}_{b}")
            # continue
        else:
            intersects: List[Polygon] = [intersection]
        for intersect in intersects:
            assert isinstance(intersect, Polygon)
            # print(intersect)
            # draw.rectangle(
            #     intersect.bounds,
            #     # fill="#000",
            #          )
            global_segmentation = list(intersect.exterior.coords)
            global_boundaries = intersect.bounds
            draw.polygon(
                global_segmentation,
                fill=(
                    255 * (category_id % 1),
                    255 * (category_id % 2),
                    255 * (category_id % 4),
                ),
                outline="#f00",
            )

            local_segmentation = [x for xs in global_segmentation for x in xs]
            local_segmentation = [[
                x - (i * (c % 2)) - (j * (1 - c % 2))
                for c, x in enumerate(local_segmentation)
            ]]
            local_boundaries = [
                global_boundaries[0] - j,
                global_boundaries[1] - i,
                global_boundaries[2] - j,
                global_boundaries[3] - i,
            ]

            annotation = dict(
                category_id=category_id,
                category_name=f"{category_name} ({kind})",
                ignore=0,
                iscrowd=0,
                bbox=local_boundaries,
                bbox_mode=BoxMode.XYXY_ABS,
                segmentation=local_segmentation,
            )
            annotations.append(annotation)

    if annotations:
        images.append(image)
    # end loop

    # print(f"Saving file {a}_{b}.tif ...")
    # img.crop((b, a, b + d, a + d)).save(filename)
# img.save("data/images/westend/belag/zzz.tif")

with open("../data/json/belaege.json", "w") as file:
    json.dump(images, file, indent=2)