import os
import json
import argparse
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

from itertools import chain
from shapely.geometry import Polygon, MultiPolygon

from PIL import Image, ImageDraw

from clean_data import CleanCollision
from data.geo_tools import calculate_piou
from data.house import SimpleHouse, plot_polys


category2color = {
    'door': '#139C5A', 'window': '#0000FF', 'bed': '#FF0000', 'cabinet': '#FFFF00',
    'bedside_table': '#F08080', 'table': '#A52A2A', 'bed_background': "#FF3333",
    'sofa': '#FF9933', 'tv_cabinet': '#FFCC99',
    'coffea_table': '#CCFF99', 'side_coffea_table': '#99FFCC', 'dining_table': '#FF6666',
    'cabinet_table': "#FF9999", 'shoe_cabinet': '#006633', 'single_sofa': '#CC6600',
    'sofa_background': '#99004C', 'cooker_cabinet': '#000099', 'sink_cabinet': '#0000CC',
    'electrical_floor_cabinet': '#3333FF', 'double_door_floor_cabinet': '#6666FF',
    'single_door_floor_cabinet': '#9999FF', 'refrigerator': '#006666', 'washbasin': '#CC0066',
    'shower': '#33FF99', 'toilet': '#660033', 'washing_set': '#FF66B2', 'washing_machine': '#FFCCE5',
    'leisure_sofa': '#666600', 'curtain_box': '#0080FF'
}
item_category2num = {
      'bed': 0, 'cabinet': 1, 'bedside_table': 2, 'table': 3, 'bed_background': 4,
      'curtain_box': 5, 'sofa': 6, 'tv_cabinet': 7,
      'coffea_table': 8, 'side_coffea_table': 9, 'dining_table': 10,
      'cabinet_table': 11, 'shoe_cabinet': 12, 'single_sofa': 13,
      'sofa_background': 14, 'cooker_cabinet': 15, 'sink_cabinet': 16,
      'electrical_floor_cabinet': 17, 'double_door_floor_cabinet': 18,
      'single_door_floor_cabinet': 19, 'refrigerator': 20, 'washbasin': 21,
      'shower': 22, 'toilet': 23, 'washing_set': 24, 'washing_machine': 25,
      'leisure_sofa': 26
}


def write_obb_text(obbs, out_file):
    text_li = []
    for obb in obbs:
        text = f"{obb['enum']} " + ' '.join([f"{p[0]} {p[1]}" for p in obb["points"]])
        text_li.append(text)
    with open(out_file, 'w') as f:
        f.write('\n'.join(text_li))


def get_obbs(poly_li, minx, miny, max_xy):
    pos_li = []
    for poly in poly_li:
        points = poly.poly.exterior.coords[:-1]
        points = [((point[0] - minx) / max_xy, (point[1] - miny) / max_xy) for point in points]
        pos_li.append(points)
    return pos_li


def get_draw_bbox(minx, miny, maxx, maxy, box_buffer=0.05):
    max_xy = max(maxx - minx, maxy - miny) * (1 + box_buffer * 2)
    x_resd = (max_xy - (maxx - minx)) / 2
    y_resd = (max_xy - (maxy - miny)) / 2
    return minx - x_resd, miny - y_resd, maxx + x_resd, maxy + y_resd


def get_center(points):
    x = sum([p[0] for p in points]) / len(points)
    y = sum([p[1] for p in points]) / len(points)
    return x, y


def draw_obbs(img_file, obb_file, out_file):
    img = Image.open(img_file)
    width, height = img.size
    draw = ImageDraw.Draw(img)
    obb_li = []
    with open(obb_file, 'r') as f:
        for line in f.readlines():
            x_li = line.split(' ')
            obb = {'enum': x_li[0],
                   'points': [(float(x_li[i]) * width, float(x_li[i+1]) * height) for i in range(1, len(x_li), 2)]}
            obb_li.append(obb)
    for obb in obb_li:
        center = get_center(obb['points'])
        draw.polygon(obb['points'], outline='yellow')
        draw.text(center, obb['enum'], fill='red', align='center')
    img.save(out_file)


def plot_layout(data, layout_img_file=None, fig_size=1024, box_buffer=0.01, alpha=0.3):
    house = SimpleHouse(data)
    minx, miny, maxx, maxy = get_draw_bbox(*house.out_room.poly.bounds, box_buffer=box_buffer)
    dpi = 72
    fig, ax = plt.subplots(figsize=(fig_size / dpi, fig_size / dpi), dpi=dpi)
    ax.set_xlim(minx, maxx)
    ax.set_ylim(miny, maxy)
    plt.axis('off')
    plt.gca().invert_yaxis()
    plt.gca().set_aspect('equal', adjustable='box')
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0, wspace=0, hspace=0)
    holes = [x.poly.exterior.coords[:-1] for x in
             chain(house.rooms.values(), house.doors.values(), house.windows.values())]
    wall = Polygon(house.out_room.poly.exterior.coords, holes=holes).buffer(0)
    colors, polys = [], []
    if isinstance(wall, MultiPolygon):
        for w in wall.geoms:
            polys.append(w)
            colors.append('#000000')
    elif isinstance(wall, Polygon):
        polys.append(wall)
        colors.append('#000000')
    else:
        print('error wall')
        return
    for door in house.doors.values():
        polys.append(door.poly)
        colors.append(category2color['door'])
    for win in house.windows.values():
        polys.append(win.poly)
        colors.append(category2color['window'])
    plot_polys(ax, polys, colors, alpha=alpha)

    f_colors = [category2color[f.category] for f in house.furnitures.values()]
    plot_polys(ax, [p.poly for p in house.furnitures.values()], f_colors, alpha=alpha)
    if layout_img_file is None:
        plt.show()
        return
    plt.savefig(layout_img_file, pad_inches=0)
    plt.close(fig)


def plot_and_save_segments(file, data_path, out_path, clean_flag=False, fig_size=1024, box_buffer=0.01, alpha=0.3, debug=False):
    layout_img_file = os.path.join(out_path, 'layout', file.replace('.json', '.png'))
    with open(os.path.join(data_path, file), 'r', encoding='utf8') as f:
        data = json.load(f)
    if clean_flag:
        ch = CleanCollision(data)
        data = ch.clean_collision()

    house = SimpleHouse(data)
    if len(house.rooms) < 1 or len(house.furnitures) < 1:
        print('no rooms or furnitures')
        return
    for room in house.rooms.values():
        if not house.out_room.poly.covers(room.poly):
            print('out room not covers room')
            return
    for x in ['frame', 'layout', 'frame_labels', 'layout_labels']:
        if not os.path.exists(os.path.join(out_path, x)):
            os.makedirs(os.path.join(out_path, x))
    minx, miny, maxx, maxy = get_draw_bbox(*house.out_room.poly.bounds, box_buffer=box_buffer)
    max_xy = max(maxx - minx, maxy - miny)
    frame_cat_obbs = []
    for obb in get_obbs(house.rooms.values(), minx, miny, max_xy):
        frame_cat_obbs.append({'enum': 0, 'points': obb})
    obb_file = os.path.join(out_path, 'frame_labels', file.replace('.json', '.txt'))
    write_obb_text(frame_cat_obbs, obb_file)
    fur_cat_obbs = []
    for obb in get_obbs(house.doors.values(), minx, miny, max_xy):
        fur_cat_obbs.append({'enum': 27, 'points': obb})
    for obb in get_obbs(house.windows.values(), minx, miny, max_xy):
        fur_cat_obbs.append({'enum': 28, 'points': obb})
    for f, obb in zip(house.furnitures.values(), get_obbs(house.furnitures.values(), minx, miny, max_xy)):
        fur_cat_obbs.append({'enum': item_category2num[f.category], 'points': obb})
    fur_obb_file = os.path.join(out_path, 'layout_labels', file.replace('.json', '.txt'))
    write_obb_text(fur_cat_obbs, fur_obb_file)
    dpi = 72
    fig, ax = plt.subplots(figsize=(fig_size / dpi, fig_size / dpi), dpi=dpi)
    ax.set_xlim(minx, maxx)
    ax.set_ylim(miny, maxy)
    plt.axis('off')
    plt.gca().invert_yaxis()
    plt.gca().set_aspect('equal', adjustable='box')
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0, wspace=0, hspace=0)
    holes = [x.poly.exterior.coords[:-1] for x in chain(house.rooms.values(), house.doors.values(), house.windows.values())]
    wall = Polygon(house.out_room.poly.exterior.coords, holes=holes).buffer(0)
    colors, polys = [], []
    if isinstance(wall, MultiPolygon):
        for w in wall.geoms:
            polys.append(w)
            colors.append('#000000')
    elif isinstance(wall, Polygon):
        polys.append(wall)
        colors.append('#000000')
    else:
        print('error wall')
        return
    for door in house.doors.values():
        polys.append(door.poly)
        colors.append(category2color['door'])
    for win in house.windows.values():
        polys.append(win.poly)
        colors.append(category2color['window'])
    plot_polys(ax, polys, colors, alpha=alpha)
    plt.savefig(os.path.join(out_path, 'frame', file.replace('.json', '.png')), pad_inches=0)
    f_colors = [category2color[f.category] for f in house.furnitures.values()]
    plot_polys(ax, [p.poly for p in house.furnitures.values()], f_colors, alpha=alpha)
    plt.savefig(layout_img_file, pad_inches=0)
    plt.close(fig)
    if debug:
        draw_obbs(os.path.join(data_path, 'frame', file.replace('.json', '.png')), obb_file,
                  os.path.join(out_path, file.replace('.json', 'debug_f.png')))
        draw_obbs(layout_img_file, fur_obb_file,
                  os.path.join(out_path, file.replace('.json', 'debug_l.png')))
    return minx, miny, maxx, maxy


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('-j', type=str, help='Input json filename')
    parser.add_argument('-d', type=str, help='Input data path')
    parser.add_argument('-o', type=str, help='Output path')
    parser.add_argument('-c', action='store_true', help='clean data')
    args = parser.parse_args()
    plot_and_save_segments(args.j, args.d, args.o, args.c, debug=False)
