import os
import shutil
import json
import tqdm
import asyncio
import aiofiles
import traceback

import pandas as pd

from data.geo_tools import get_door_buff_poly, calculate_piou, calculate_door_iod
from data.house import SimpleHouse


room_map = {1: 'living', 3: 'living', 4: 'bedroom', 5: 'bedroom', 6: 'bedroom', 7: 'kitchen',
            8: 'bathroom', 9: 'balcony', 10: 'bathroom', 20: 'balcony', 21: 'balcony', 27: 'storage',
            30: 'others', 32: 'bathroom', 34: 'bathroom', 35: 'balcony'}
item_category_li = ['cabinet', 'bedside_table', 'table', 'bed_background', 'curtain_box', 'sofa', 'tv_cabinet',
                    'coffea_table', 'side_coffea_table', 'dining_table', 'cabinet_table', 'shoe_cabinet', 'bed',
                    'single_sofa', 'sofa_background', 'cooker_cabinet', 'sink_cabinet', 'electrical_floor_cabinet',
                    'double_door_floor_cabinet', 'single_door_floor_cabinet', 'refrigerator', 'washbasin', 'shower',
                    'toilet', 'washing_set', 'washing_machine', 'leisure_sofa']


def layout_statistics(data):
    res = {k: 0 for k in item_category_li}
    h = SimpleHouse(data)
    res['empty_room_num'] = 0
    res['room_num'] = len(h.rooms)
    for rid in h.rooms:
        if len(h.room2furnitures.get(rid, [])) < 1:
            res['empty_room_num'] += 1
    for furniture in data['furniture']:
        res[furniture['type']] += 1
    furniture_polys = [f.poly for f in h.furnitures.values() if f.category != 'curtain_box']
    door_polys = [get_door_buff_poly(d.poly) for d in h.doors.values()]
    piou_li = calculate_piou(furniture_polys)
    res['furniture_pair_num'] = len(piou_li)
    res['furniture_inter_num'] = len([p for p in piou_li if p > 1e-4])
    res['furniture_piou_sum'] = sum(piou_li)
    door_iod_list = calculate_door_iod(door_polys, furniture_polys)
    res['door_num'] = len(door_iod_list)
    res['door_block_num'] = len([d for d in door_iod_list if d > 0])
    res['door_iod_sum'] = sum(door_iod_list)
    return res


async def read_json(file_name):
    async with aiofiles.open(file_name, mode='r', encoding='utf8') as f:
        data = await f.read()
    return json.loads(data)


async def single_run(filename, input_path):
    try:
        plan_data = await read_json(os.path.join(input_path, filename))
        return layout_statistics(plan_data)
    except:
        print(filename, traceback.format_exc())
        return None


async def batch_run(filenames, input_path):
    tasks = [asyncio.create_task(single_run(f, input_path)) for f in filenames]
    results = await asyncio.gather(*tasks)
    df = pd.DataFrame(results)
    filename = 'ours_statistics.csv'
    if os.path.exists(filename):
        df.to_csv(filename, mode='a', header=False, index=False)
    else:
        df.to_csv(filename, index=False)
    return results


if __name__ == '__main__':
    input_path = '/anonymous/vectors_in_one'
    chunk_size = 10
    filenames = []
    for filename in os.listdir(input_path):
        if filename.endswith('.json'):
            filenames.append(filename)
    filenames = filenames[:10000]
    for i in tqdm.tqdm(range(0, len(filenames), chunk_size)):
        asyncio.get_event_loop().run_until_complete(batch_run(filenames[i: i+chunk_size], input_path))
