import os
import math
import json
import tqdm
import traceback

import pandas as pd
import matplotlib.pyplot as plt

from shapely.geometry import Polygon, LineString
from shapely.ops import unary_union

# https://github.com/3D-FRONT-FUTURE/3D-FRONT-ToolBox
# import from 3D-FRONT-ToolBox, copy this file to 3D-FRONT-ToolBox project root directory
from Room.Preprocess.src.floorplan_generator import FloorplanGenerator
from Generator.model import ModelInstance


category_map = {
    "Nightstand": "bedside_table", "Wardrobe": "cabinet", "Three-Seat / Multi-seat Sofa": "sofa",
    "Dining Table": "dining_table", "Coffee Table": "coffea_table", "Loveseat Sofa": "sofa",
    "Children Cabinet": "cabinet", "Drawer Chest / Corner cabinet": "cabinet", "King-size Bed": "bed",
    "TV Stand": "tv_cabinet", "Sideboard / Side Cabinet / Console": "cabinet_table", "Lazy Sofa": "leisure_sofa",
    "Dressing Table": "table", "Wine Cabinet": "cabinet_table", "L-shaped Sofa": "sofa", "Corner/Side Table":
        "side_coffea_table", "Bookcase / jewelry Armoire": "cabinet", "Kids Bed": "bed",
    "Sideboard / Side Cabinet / Console Table": "table", "Bed Frame": "bed", "Shoe Cabinet": "shoe_cabinet",
    "Three-Seat / Multi-person sofa": "sofa", "Double Bed": "bed", "Bunk Bed": "bed", "Desk": "table",
    "Two-seat Sofa": "sofa", "Tea Table": "coffea_table", "Couch Bed": "bed", "Single bed": "bed",
    "Chaise Longue Sofa": "sofa", "U-shaped Sofa": "sofa"
}


class FloorFurniture(FloorplanGenerator):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.furnitures = []

    def generate_furniture_info(self, furniture_info):
        furniture_dict = {x['uid']: x for x in self.json_information_acquirer.json_data['furniture']}
        for room in self.json_information_acquirer.json_data['scene']['room']:
            for children in room['children']:
                if any([math.isnan(x) for x in children['pos']]) \
                        or children['ref'] not in furniture_dict:
                    continue
                jid = furniture_dict[children['ref']]['jid']
                if jid not in furniture_info:
                    continue
                model_info = furniture_info[jid]
                data = {**children, **model_info}
                data['jid'] = jid
                data['size'] = data['boundingBox']
                model_inst = ModelInstance(data)
                model_inst.category = model_info['category']
                model_inst.super_category = model_info['super-category']
                self.furnitures.append(model_inst)


def get_points_from_list(xy_list, scale=1):
    points = []
    for i in range(0, len(xy_list), 2):
        points.append([scale * xy_list[i], scale * xy_list[i + 1]])
    return points


def plot_floor(floor_inst, save_file=None):
    fig, ax = plt.subplots(figsize=(26, 14), dpi=72)
    polys, labels = [], []
    floor_dict = floor_inst.house_info_dict
    for room in floor_dict['floorplan']:
        polys.append(Polygon(get_points_from_list(room['floor'])))
        labels.append(room['room'])
        for key in ['door', 'window', 'baywindow', 'hole']:
            for door in room[key]:
                polys.append(Polygon(get_points_from_list(door)))
                labels.append(key)
    for poly, label in zip(polys, labels):
        center = poly.centroid
        ax.plot(*poly.exterior.coords.xy)
        ax.text(center.x, center.y, label, horizontalalignment='center', verticalalignment='center')
    for furn in floor_inst.furnitures:
        coords = furn.get_bounding_box().tolist()
        poly = Polygon(coords)
        center = poly.centroid
        ax.plot(*poly.exterior.coords.xy)
        ax.text(center.x, center.y, furn.category, horizontalalignment='center', verticalalignment='center')
    plt.gca().invert_yaxis()
    plt.gca().set_aspect('equal', adjustable='box')
    if save_file is None:
        plt.show()
    else:
        plt.savefig(save_file)


def direction2angle(direction):
    return round(math.degrees(math.atan2(direction[1], direction[0])), 2)


def get_vector(wall):
    start, end = wall
    length = math.sqrt(sum([(end[i] - start[i]) ** 2 for i in range(2)]))
    if length < 1e-6:
        raise Exception('zero length line, coords: {0}, {1}'.format(start, end))
    vector = [(end[i] - start[i]) / length for i in range(2)]
    return vector


def get_lines_from_points(points):
    lines = []
    for i in range(len(points)):
        lines.append(LineString([points[i], points[(i + 1) % len(points)]]))
    return lines


def get_out_room(room_polys, win_door_polys):
    for buffer_distance in range(10, 30, 1):
        out_room = unary_union([room.buffer(buffer_distance, cap_style=3, join_style=2) for room in room_polys])
        if isinstance(out_room, Polygon) and all([out_room.covers(x) for x in win_door_polys]):
            room_data = {'roomId': 'out', 'roomName': 'out', 'roomType': 0, 'wallPoints': out_room.exterior.coords[:-1]}
            return room_data


def convert2simple_json(floor_inst, save_file):
    simple_json = {'rooms': [], 'windowsDoors': [], 'furniture': []}
    floor_dict = floor_inst.house_info_dict
    room_polys, win_door_polys = [], []
    for i, room in enumerate(floor_dict['floorplan']):
        coords = get_points_from_list(room['floor'], scale=100)[:-1]
        room_polys.append(Polygon(coords))
        simple_json['rooms'].append(
            {'roomId': f'room_{i}', 'roomName': room['room'], 'roomType': 1,
             'wallPoints': coords})
        for key in ['door', 'window', 'baywindow', 'hole']:
            for door_idx, door in enumerate(room[key]):
                door_coords = get_points_from_list(door, scale=100)
                door_lines = get_lines_from_points(door_coords)
                door_lines.sort(key=lambda x: x.length)
                angle = direction2angle(get_vector(door_lines[-1].coords[:]))
                door_poly = Polygon(door_coords)
                pos = door_poly.centroid
                win_door_polys.append(door_poly)
                height2floor = room[f'{key}_height'] if isinstance(room[f'{key}_height'], float) \
                    else room[f'{key}_height'][door_idx]
                height = room[f'{key}_height_limit'] if isinstance(room[f'{key}_height_limit'], float) \
                    else room[f'{key}_height_limit'][door_idx]
                simple_json['windowsDoors'].append(
                    {'type': 'door' if key in ['door', 'hole'] else 'window',
                     'pos': (pos.x, pos.y, height2floor), 'rotate': angle,
                     'length': door_lines[-1].length, 'width': door_lines[0].length, 'height': height})
    simple_json['rooms'].append(get_out_room(room_polys, win_door_polys))
    for furn in floor_inst.furnitures:
        if furn.category not in category_map:
            continue
        pos = (100 * furn.pos).tolist()
        simple_json['furniture'].append(
            {'type': category_map[furn.category], 'pos': (pos[0], pos[2], pos[1]),
             'length': 100 * furn.size[0], 'width': 100 * furn.size[2],
             'height': 100 * furn.size[1], 'rotate': direction2angle(furn.dir)-90})
    with open(save_file, 'w') as f:
        json.dump(simple_json, f)


def get_model_info(json_file, csv_file):
    with open(json_file, 'r') as f:
        model_info = json.load(f)
    category_df = pd.read_csv(csv_file)
    for data in category_df.to_dict(orient='records'):
        if data['model_id'] in model_info:
            model_info[data['model_id']].update({k: data[k] for k in ['category', 'super-category']})
    return model_info


if __name__ == '__main__':
    scene_dir = 'D:\\data\\3D-FRONT'
    simple_json_path = 'D:\\data\\3D-FRONT-SIMPLE'
    if not os.path.exists(simple_json_path):
        os.makedirs(simple_json_path)
    model_category_file = 'scripts\\scene_synthesis\\deep_synth\\data\\ModelCategoryMapping.csv'

    model_info = get_model_info('config/model_info.json', model_category_file)
    for file in tqdm.tqdm(os.listdir(scene_dir)):
        if not file.endswith('.json') or os.path.exists(os.path.join(simple_json_path, file)):
            continue
        try:
            inst = FloorFurniture()
            inst.generate_floorplan(os.path.join(scene_dir, file))
            inst.generate_furniture_info(model_info)
            convert2simple_json(inst, os.path.join(simple_json_path, file))
        except Exception:
            print('failed to process file:', file)
            traceback.print_exc()
