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


from tools.generate_text import SimpleHouse, plot_polys
from data_driven.clean_data import CleanCollision
from .layout_single_gen import category2color


def plot_split_rooms(file, data_path, out_path, clean_flag, fig_size=256, alpha=0.3, debug=False):
    with open(os.path.join(data_path, file), 'r', encoding='utf8') as f:
        data = json.load(f)

    house = SimpleHouse(data)
    if len(house.rooms) < 1 or len(house.furnitures) < 1:
        print('no rooms or furnitures')
        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))

    dpi = 72
    for rid, room in house.rooms.items():
        room_type = None
        if 'bedroom' in room.name.lower():
            room_type = 'bedroom'
        elif 'living' in room.name.lower():
            room_type = 'living_room'
        else:
            continue
        fig, ax = plt.subplots(figsize=(fig_size / dpi, fig_size / dpi), dpi=dpi)
        plt.axis('off')
        plt.gca().invert_yaxis()
        plt.gca().set_aspect('equal', adjustable='box')
        polys = [house.furnitures[fid] for fid in house.room2furnitures[rid]]
        if len(polys) < 1:
            continue
        colors = [category2color[f.category] for f in polys]
        plot_polys(ax, polys, colors, alpha=alpha)
        if not os.path.exists(os.path.join(out_path, 'rooms', room_type)):
            os.makedirs(os.path.join(out_path, 'rooms', room_type))
        plt.savefig(os.path.join(out_path, 'rooms', room_type, f"{room.name}.png"), pad_inches=0)
        plt.close(fig)


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', type=bool, default=False, help='clean data')
    args = parser.parse_args()
    plot_split_rooms(args.j, args.d, args.o, args.c, debug=False)
