"""
This code is an adaptation that uses Structured 3D for the code base.

Reference: https://github.com/bertjiazheng/Structured3D
"""

import json
import os
import sys

import numpy as np
from shapely.geometry import Polygon

sys.path.append("../data_preprocess")
from common_utils import resort_corners

type2id = {
    "living room": 0,
    "kitchen": 1,
    "bedroom": 2,
    "bathroom": 3,
    "balcony": 4,
    "corridor": 5,
    "dining room": 6,
    "study": 7,
    "studio": 8,
    "store room": 9,
    "garden": 10,
    "laundry room": 11,
    "office": 12,
    "basement": 13,
    "garage": 14,
    "undefined": 15,
    "door": 16,
    "window": 17,
}


def generate_density(point_cloud, width=256, height=256):

    ps = point_cloud * -1
    ps[:, 0] *= -1
    ps[:, 1] *= -1

    image_res = np.array((width, height))

    max_coords = np.max(ps, axis=0)
    min_coords = np.min(ps, axis=0)
    max_m_min = max_coords - min_coords

    max_coords = max_coords + 0.1 * max_m_min
    min_coords = min_coords - 0.1 * max_m_min

    normalization_dict = {}
    normalization_dict["min_coords"] = min_coords
    normalization_dict["max_coords"] = max_coords
    normalization_dict["image_res"] = image_res

    # coordinates = np.round(points[:, :2] / max_coordinates[None,:2] * image_res[None])
    coordinates = np.round(
        (ps[:, :2] - min_coords[None, :2]) / (max_coords[None, :2] - min_coords[None, :2]) * image_res[None]
    )
    coordinates = np.minimum(np.maximum(coordinates, np.zeros_like(image_res)), image_res - 1)

    density = np.zeros((height, width), dtype=np.float32)

    unique_coordinates, counts = np.unique(coordinates, return_counts=True, axis=0)
    # print(np.unique(counts))
    # counts = np.minimum(counts, 1e2)

    unique_coordinates = unique_coordinates.astype(np.int32)

    density[unique_coordinates[:, 1], unique_coordinates[:, 0]] = counts
    density = density / np.max(density)

    return density, normalization_dict


def normalize_point(point, normalization_dict):

    min_coords = normalization_dict["min_coords"]
    max_coords = normalization_dict["max_coords"]
    image_res = normalization_dict["image_res"]

    point_2d = np.round((point[:2] - min_coords[:2]) / (max_coords[:2] - min_coords[:2]) * image_res)
    point_2d = np.minimum(np.maximum(point_2d, np.zeros_like(image_res)), image_res - 1)

    point[:2] = point_2d.tolist()

    return point


def normalize_annotations(scene_path, normalization_dict):
    annotation_path = os.path.join(scene_path, "annotation_3d.json")
    with open(annotation_path, "r") as f:
        annotation_json = json.load(f)

    for line in annotation_json["lines"]:
        point = line["point"]
        point = normalize_point(point, normalization_dict)
        line["point"] = point

    for junction in annotation_json["junctions"]:
        point = junction["coordinate"]
        point = normalize_point(point, normalization_dict)
        junction["coordinate"] = point

    return annotation_json


def parse_floor_plan_polys(annos):
    planes = []
    for semantic in annos["semantics"]:
        for planeID in semantic["planeID"]:
            if annos["planes"][planeID]["type"] == "floor":
                planes.append({"planeID": planeID, "type": semantic["type"]})

        # if semantic["type"] == "outwall":
        #     outerwall_planes = semantic["planeID"]

    # extract hole vertices
    lines_holes = []
    for semantic in annos["semantics"]:
        if semantic["type"] in ["window", "door"]:
            for planeID in semantic["planeID"]:
                lines_holes.extend(np.where(np.array(annos["planeLineMatrix"][planeID]))[0].tolist())
    lines_holes = np.unique(lines_holes)

    ## junctions on the floor
    # junctions = np.array([junc["coordinate"] for junc in annos["junctions"]])

    # construct each polygon
    polygons = []
    for plane in planes:
        lineIDs = np.where(np.array(annos["planeLineMatrix"][plane["planeID"]]))[0].tolist()
        junction_pairs = [np.where(np.array(annos["lineJunctionMatrix"][lineID]))[0].tolist() for lineID in lineIDs]
        polygon = convert_lines_to_vertices(junction_pairs)
        polygons.append([polygon[0], plane["type"]])

    return polygons


def convert_lines_to_vertices(lines):
    """
    convert line representation to polygon vertices

    """
    polygons = []
    lines = np.array(lines)

    polygon = None
    while len(lines) != 0:
        if polygon is None:
            polygon = lines[0].tolist()
            lines = np.delete(lines, 0, 0)

        lineID, juncID = np.where(lines == polygon[-1])
        vertex = lines[lineID[0], 1 - juncID[0]]
        lines = np.delete(lines, lineID, 0)

        if vertex in polygon:
            polygons.append(polygon)
            polygon = None
        else:
            polygon.append(vertex)

    return polygons


def generate_coco_dict(annos, polygons, curr_instance_id, curr_img_id, ignore_types):

    junctions = np.array([junc["coordinate"][:2] for junc in annos["junctions"]])

    coco_annotation_dict_list = []

    for poly_ind, (polygon, poly_type) in enumerate(polygons):
        if poly_type in ignore_types:
            continue

        polygon = junctions[np.array(polygon)]

        poly_shapely = Polygon(polygon)
        area = poly_shapely.area

        # assert area > 10
        # if area < 100:
        if poly_type not in ["door", "window"] and area < 100:
            continue
        if poly_type in ["door", "window"] and area < 1:
            continue

        rectangle_shapely = poly_shapely.envelope

        ### here we convert door/window annotation into a single line
        if poly_type in ["door", "window"]:
            assert polygon.shape[0] == 4
            midp_1 = (polygon[0] + polygon[1]) / 2
            midp_2 = (polygon[1] + polygon[2]) / 2
            midp_3 = (polygon[2] + polygon[3]) / 2
            midp_4 = (polygon[3] + polygon[0]) / 2

            dist_1_3 = np.square(midp_1 - midp_3).sum()
            dist_2_4 = np.square(midp_2 - midp_4).sum()
            if dist_1_3 > dist_2_4:
                polygon = np.row_stack([midp_1, midp_3])
            else:
                polygon = np.row_stack([midp_2, midp_4])

        coco_seg_poly = []
        poly_sorted = resort_corners(polygon)

        for p in poly_sorted:
            coco_seg_poly += list(p)

        # Slightly wider bounding box
        bound_pad = 2
        bb_x, bb_y = rectangle_shapely.exterior.xy
        bb_x = np.unique(bb_x)
        bb_y = np.unique(bb_y)
        bb_x_min = np.maximum(np.min(bb_x) - bound_pad, 0)
        bb_y_min = np.maximum(np.min(bb_y) - bound_pad, 0)

        bb_x_max = np.minimum(np.max(bb_x) + bound_pad, 256 - 1)
        bb_y_max = np.minimum(np.max(bb_y) + bound_pad, 256 - 1)

        bb_width = bb_x_max - bb_x_min
        bb_height = bb_y_max - bb_y_min

        coco_bb = [bb_x_min, bb_y_min, bb_width, bb_height]

        coco_annotation_dict = {
            "segmentation": [coco_seg_poly],
            "area": area,
            "iscrowd": 0,
            "image_id": curr_img_id,
            "bbox": coco_bb,
            "category_id": type2id[poly_type],
            "id": curr_instance_id,
        }

        coco_annotation_dict_list.append(coco_annotation_dict)
        curr_instance_id += 1

    return coco_annotation_dict_list