import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import tqdm
import uuid
import itertools
import torch
import pickle
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

from cleanfid import fid
from PIL import Image
from diffusers import DPMSolverSDEScheduler
from models.unet import FrameConditionalUNet2DModel
from pipelines.frame_cond_pipeline import FrameConditionalDDPMPipeline

from shapely.geometry import Polygon
from ultralytics import YOLO
from post_process import post_process_box_data, replace_color
from layout_single_gen import category2color
from data.house import plot_polys



img_size = (256, 256)
num_inference_steps = 100
model_path = 'D:\\trains\\fp_cond_09111518'
# model_path = 'D:\\trains\\frame_cond_al_clean'
# frame_data_path = 'D:\\data\\layouts\\cond\\test'
# frame_data_path = 'D:\\data\\3D-FRONT-SIMPLE-CLEAN\\frame'
frame_data_path = 'D:\\data\\fp_json_split\\frame\\train'
layout_data_path = 'D:\\data\\fp_json_split\\layout\\train'


# frame_label_path = 'D:\\data\\3D-FRONT-SIMPLE-CLEAN\\frame_labels'
frame_label_path = 'D:\\data\\fp_json_split\\frame_labels\\train'
# output_path = 'D:\\data\\frame_cond_test_0924'
output_path = 'D:\\data\\frame_cond_al_test_0926'
output_path_rs = 'D:\\data\\frame_cond_al_rs_test_0926'
# output_path_rs = 'D:\\data\\frame_cond_rs_test_0924x'
# output_path = 'D:\\data\\frame_cond_test_0912'
device_name = 'cuda:1'
device = torch.device(device_name)
model = FrameConditionalUNet2DModel.from_pretrained(model_path, subfolder='unet')
model = model.to(device)
noise_scheduler = DPMSolverSDEScheduler.from_pretrained(model_path, subfolder='scheduler')
pipe = FrameConditionalDDPMPipeline(model, noise_scheduler)
pipe = pipe.to(device)
generator = torch.Generator(device=device).manual_seed(0)
obb_model_path = 'D:\\trains\\obbs\\best0912.pt'
# obb_model_path = 'D:\\trains\\obbs\\best_ali.pt'
obb_model = YOLO(obb_model_path)


def calculate_pioc(polys):
    if len(polys) < 2:
        return []
    piou_li = []
    for pa, pb in itertools.combinations(polys, 2):
        if pa.intersects(pb):
            current_val = pa.intersection(pb).area / pa.union(pb).area
            piou_li.append(current_val)
        else:
            piou_li.append(0)
    return piou_li


def read_polygons_from_file(file_path):
    polygons = []
    with open(file_path, 'r') as file:
        for line in file:
            # 忽略第一列，并将剩余的数据转换为浮点数
            data = list(map(float, line.split()[1:]))
            # 将数据对分为坐标对
            coordinates = [(data[i], data[i + 1]) for i in range(0, len(data), 2)]
            # 创建多边形
            polygon = Polygon(coordinates)
            polygons.append(polygon)
    return polygons


def get_poly_group_li(room_polys, box_data):
    furniture_polys = [{'polys': [], 'poly_cat_li': [], 'room_type': ''} for _ in range(len(room_polys))]
    all_polys = []
    for furniture in box_data:
        if furniture['label'] in ['door', 'window']:
            continue
        poly = Polygon(furniture['points'])
        all_polys.append(poly)
        for room_idx, room in enumerate(room_polys):
            if room.covers(poly.centroid):
                furniture_polys[room_idx]['polys'].append(poly)
                furniture_polys[room_idx]['poly_cat_li'].append(furniture['label'])
                break
    for fp in furniture_polys:
        if 'bed' in fp['poly_cat_li']:
            fp['room_type'] = 'bedroom'
        elif 'sofa' in fp['poly_cat_li'] or 'tv_cabinet' in fp['poly_cat_li']:
            fp['room_type'] = 'living_room'

    return furniture_polys, all_polys


def get_filename_without_extension(file_path):
    # 获取文件名和扩展名
    file_name, file_extension = os.path.splitext(os.path.basename(file_path))
    return file_name


def plot_furniture_without_room(box_data, frame_file, out_path):
    fig_size = 512
    dpi = 72
    room_polys = read_polygons_from_file(frame_file)
    poly_group_li, all_polys = get_poly_group_li(room_polys, box_data)
    uid = get_filename_without_extension(frame_file)
    living_piou_li = []
    bedroom_piou_li = []
    for idx, group in enumerate(poly_group_li):
        if group["room_type"] not in ['bedroom', 'living_room']:
            continue
        filename = f'{uid}_{idx}.png'
        file_path = os.path.join(out_path, group["room_type"])
        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')
        plot_polys(ax, group['polys'], [category2color[c] for c in group['poly_cat_li']], alpha=0.3)
        if not os.path.exists(file_path):
            os.makedirs(file_path)
        plt.savefig(os.path.join(file_path, filename))
        plt.close()
        pious = calculate_pioc(group['polys'])
        if group["room_type"] == 'bedroom':
            bedroom_piou_li.extend(pious)
        elif group['room_type'] == 'living_room':
            living_piou_li.extend(pious)
    return living_piou_li, bedroom_piou_li

batch_size = 4
if not os.path.exists(output_path):
    os.makedirs(output_path)
if not os.path.exists(output_path_rs):
    os.makedirs(output_path_rs)

files = [f for f in os.listdir(frame_data_path) if f.endswith('png')][1000:2000]
for i in tqdm.tqdm(range(0, len(files), batch_size)):
    conds = []
    for img_file in files[i: i+batch_size]:
        img = Image.open(os.path.join(frame_data_path, img_file)).resize(img_size).convert('RGB')
        cond = torch.from_numpy(np.array(img)).float()
        conds.append(cond.permute(2, 0, 1))
    c_batch_size = len(conds)
    conds = torch.stack(conds).to(device)
    pred_images = pipe(conds, batch_size=c_batch_size, generator=generator, num_inference_steps=num_inference_steps)
    for j, p_img in enumerate(pred_images):
        p_img.save(os.path.join(output_path, files[i+j]))

score = fid.compute_fid(layout_data_path, output_path, mode="clean", num_workers=0)
print('fid', score)
score = fid.compute_kid(layout_data_path, output_path, num_workers=0)
print('kid', score)


piou_li = []
all_living_piou_li = []
all_bedroom_piou_li = []
pred_files = [f for f in os.listdir(output_path) if f.endswith('png')]
for i in tqdm.tqdm(range(0, len(pred_files), batch_size * 10)):
    pred_images = [replace_color(Image.open(os.path.join(output_path, f)).convert('RGB'))
                   for f in pred_files[i: i+batch_size * 10]]
    # pred_images = [Image.open(os.path.join(output_path, f)).convert('RGB')
    #                for f in pred_files[i: i + batch_size * 10]]
    obb_res_li = obb_model.predict(source=pred_images, conf=0.8, iou=0.7, imgsz=256, device=device_name)
    for j, obb_res in enumerate(obb_res_li):
        box_li = post_process_box_data(obb_res)
        living_piou_li, bedroom_piou_li = plot_furniture_without_room(box_li,
                                    os.path.join(frame_label_path, f"{os.path.splitext(pred_files[i + j])[0]}.txt"),
                                    output_path_rs
                                    )
        box_li = [box for box in box_li if box['label'] not in ['curtain_box', 'door', 'window']]
        piou = calculate_pioc([Polygon(box['points']) for box in box_li])
        piou_li.extend(piou)
        all_living_piou_li.extend(living_piou_li)
        all_bedroom_piou_li.extend(bedroom_piou_li)

with open('piou_li.pkl', 'wb') as f:
    pickle.dump(piou_li, f)
# with open('piou_li.pkl', 'rb') as f:
#     piou_li = pickle.load(f)
print(np.mean(piou_li), len(piou_li))
print(len([p for p in piou_li if p > 0]) / len(piou_li))

print('living room')
with open('l_piou_li.pkl', 'wb') as f:
    pickle.dump(all_living_piou_li, f)
# with open('l_piou_li.pkl', 'rb') as f:
#     all_living_piou_li = pickle.load(f)
print(np.mean(all_living_piou_li), len(all_living_piou_li))
print(len([p for p in all_living_piou_li if p > 0]) / len(all_living_piou_li))

print('bedroom')
with open('b_piou_li.pkl', 'wb') as f:
    pickle.dump(all_bedroom_piou_li, f)
# with open('b_piou_li.pkl', 'rb') as f:
#     all_bedroom_piou_li = pickle.load(f)
print(np.mean(all_bedroom_piou_li), len(all_bedroom_piou_li))
print(len([p for p in all_bedroom_piou_li if p > 0]) / len(all_bedroom_piou_li))
