import json
import math
import numpy as np
import matplotlib.pyplot as plt

from PIL import Image
from shapely.geometry import Polygon, Point, LineString
from shapely.ops import unary_union
from shapely.affinity import affine_transform
from data.geo_tools import plot_polys, clean_polygon_colliner_point


def get_box_size(points):
    lines = []
    for i in range(1, len(points)):
        lines.append(LineString([points[i - 1], points[i]]))
    length = [x.length for x in lines]
    return min(length), max(length)


def get_draw_scale(data):
    constant_width = {'window': 24, 'cabinet': 55}
    x_li, y_li = [], []
    for d in data:
        if d['label'] in constant_width:
            w, _ = get_box_size(d['points'])
            x_li.append(w)
            y_li.append(constant_width[d['label']])
    # least square method
    x = np.array(x_li)
    y = np.array(y_li)
    return np.mean(x * y) / np.mean(x ** 2)


def get_line_vector(line):
    start, end = line.coords[0], line.coords[-1]
    assert line.length > 0
    vector = [(end[i] - start[i]) / line.length for i in range(2)]
    return vector


def get_included_cos(v1, v2):
    v1 = np.array(v1)
    v2 = np.array(v2)
    l1 = np.linalg.norm(v1)
    l2 = np.linalg.norm(v2)
    return v1.dot(v2) / (l1 * l2)


def get_line_project(l1, l2):
    proj_dis = [l2.project(Point(c)) for c in l1.coords]
    return LineString([l2.interpolate(proj_d) for proj_d in proj_dis])


def post_process_door(door, rooms, room_walls, i):
    door_poly = Polygon(door['points'])
    distances = [(door_poly.distance(rooms[i]), i) for i in range(len(rooms))]
    distances.sort(key=lambda x: x[0])
    new_lines = []
    door_lines = get_poly_lines(door_poly)
    door_length_line = max(door_lines, key=lambda x: x.length)
    door_vector = get_line_vector(door_length_line)
    for distance, room_idx in distances:
        if distance > 20:
            break
        min_distance = float('inf')
        min_wall_idx = None
        for wall_idx, (room_line, line_dir, inner_dir) in enumerate(room_walls[room_idx]):
            if room_line.length < door_length_line.length:
                continue
            current_distance = room_line.distance(door_poly)
            if current_distance < min_distance and 1 - abs(get_included_cos(door_vector, line_dir)) < 0.2:
                min_distance = current_distance
                min_wall_idx = wall_idx
        if min_wall_idx is not None:
            new_lines.append((room_walls[room_idx][min_wall_idx][0], room_walls[room_idx][min_wall_idx][2]))
    door_door_lines = []
    for nl, nv in new_lines:
        door_door_lines.append((get_line_project(door_length_line, nl), nv))
    if len(door_door_lines) >= 2:
        print('fix 2 rooms', i)
        return unary_union([x[0] for x in door_door_lines]).convex_hull, door_door_lines[0][1]
    if len(door_door_lines) == 1:
        print('fix 1 rooms', i)
        pos_diff = [-24 * door_door_lines[0][1][i] for i in range(2)]
        matrix = [1, 0, 0, 1] + pos_diff
        outer_line = affine_transform(door_door_lines[0][0], matrix)
        return unary_union([door_door_lines[0][0], outer_line]).convex_hull, door_door_lines[0][1]
    print('fix 0 rooms', i)
    return door_poly, [0, 1]


short_edge_first_li = ['bed', 'toilet']


def post_procss_furnitures(frame_data, furnitures, bounds):
    minx, miny, maxx, maxy = bounds
    maxv = max(maxx - minx, maxy - miny)
    rooms, room_walls = [], []
    for room_data in frame_data['rooms']:
        if room_data['roomType'] <= 0:
            continue
        room_poly = Polygon(room_data['wallPoints'])
        rooms.append(room_poly)
        walls = get_poly_lines(room_poly)
        wall_directions = [get_line_vector(x) for x in walls]
        walls = [(w, d, (d[1], -d[0])) for w, d in zip(walls, wall_directions)]
        room_walls.append(walls)
    # empty furniture list
    frame_data['furniture'] = []
    for furniture in furnitures:
        if furniture['label'] in ['door', 'window']:
            continue
        points = furniture['points']
        points = [[p[0] * maxv + minx, p[1] * maxv + miny] for p in points]
        poly = Polygon(points)
        short_edge_first = furniture['label'] in short_edge_first_li
        length_line, distance, direction = get_furniture_direction(poly, rooms, room_walls, thres=30,
                                                                   short_edge_first=short_edge_first)
        if length_line is not None:
            format_result(poly, furniture['label'], frame_data, length=length_line.length, direction=direction)
    return frame_data


def post_process_frame_data(data, debug=False, save_file=None):
    pretty_res = {'rooms': [], 'windowsDoors': [], 'furniture': []}
    scale = get_draw_scale(data)
    polys, labels, rooms = [], [], []
    for i, x in enumerate(data):
        if x['label'] != 'room':
            continue
        x['points'] = [[p[0] * scale, p[1] * scale] for p in x['points']]
        if len(x['points']) > 4:
            points = remove_short_edges(x['points'], 15)
            if len(points) > 3:
                poly = Polygon(points).buffer(0)
                format_result(poly, 'room', pretty_res, roomType=1, roomId=f'room_{i}')
        else:
            poly = Polygon(x['points'])
        rooms.append(poly)
        polys.append(poly)
        labels.append(f'room_{i}')
    room_walls = []
    for room in rooms:
        walls = get_poly_lines(room)
        wall_directions = [get_line_vector(x) for x in walls]
        walls = [(w, d, (d[1], -d[0])) for w, d in zip(walls, wall_directions)]
        room_walls.append(walls)
        # for w in walls:
        #     polys.append(w[0])
        #     labels.append(f'{round(w[2][0], 2)} {round(w[2][1], 2)}')

    for i, x in enumerate(data):
        if x['label'] == 'room':
            continue
        x['points'] = [[p[0] * scale, p[1] * scale] for p in x['points']]
        if x['label'] in ['door', 'window']:
            poly, direction = post_process_door(x, rooms, room_walls, i)
            if isinstance(poly, Polygon):
                format_result(poly, x['label'], pretty_res, direction=direction)
        else:
            poly = Polygon(x['points'])
            short_edge_first = x['label'] in short_edge_first_li
            length_line, distance, direction = get_furniture_direction(poly, rooms, room_walls,
                                                                       short_edge_first=short_edge_first)
            if length_line is not None:
                format_result(poly, x['label'], pretty_res, length=length_line.length, direction=direction)
        # poly = Polygon(x['points'])
        polys.append(poly)
        labels.append(f"{x['label']}_{i}")
    out_room = unary_union([room.buffer(24, cap_style=3, join_style=2) for room in rooms])
    out_room = Polygon(remove_duplicate_points(out_room.exterior.coords[:-1]))
    polys.append(out_room)
    format_result(out_room, 'room', pretty_res, roomType=0, roomId='outRoom')
    if debug or save_file:
        plot_layout(polys, labels, save_file=save_file)
    return pretty_res


def remove_duplicate_points(coords):
    seen = set()
    new_coords = []
    for coord in coords:
        coord = (round(coord[0], 2), round(coord[1], 2))
        if coord not in seen:
            seen.add(coord)
            new_coords.append(coord)
    return new_coords


def get_furniture_direction(poly, rooms, room_walls, thres=12, short_edge_first=False):
    poly_lines = get_poly_lines(poly)
    for i in range(len(rooms)):
        if rooms[i].covers(poly.centroid):
            candidate_dirs = []
            for wall, _, inner_dir in room_walls[i]:
                if wall.distance(poly) <= thres:
                    project_lines = [(get_line_project(pl, wall), pl.distance(wall)) for pl in poly_lines]
                    project_lines = [pl for pl in project_lines if (not pl[0].is_empty) and (pl[0].length > 0)]
                    if len(project_lines) == 0:
                        continue
                    project_length = max(project_lines, key=lambda x: (x[0].length, -x[1]))
                    candidate_dirs.append((*project_length, inner_dir))
            if len(candidate_dirs) and short_edge_first:
                return max(candidate_dirs, key=lambda x: (-x[0].length, x[1]))
            elif len(candidate_dirs):
                return max(candidate_dirs, key=lambda x: (x[0].length, x[1]))
    if short_edge_first:
        final_line = min(poly_lines, key=lambda x: x.length)
    else:
        final_line = max(poly_lines, key=lambda x: x.length)
    line_dir = get_line_vector(final_line)
    return final_line, 0, (line_dir[1], -line_dir[0])


def format_result(poly, category, pretty_res, length=None, direction=None, roomType=1, roomId=None):
    if category == 'room':
        data = {'roomName': '', 'roomType': roomType,
                'wallPoints': [[round(x[0], 2), round(x[1], 2)]
                               for x in clean_polygon_colliner_point(poly.exterior.coords[:-1])]}
        if roomId:
            data['roomId'] = roomId
        pretty_res['rooms'].append(data)
    else:
        pos = poly.centroid.coords[0]
        if length is None:
            width, length = get_box_size(poly.exterior.coords[:-1])
        else:
            s, b = get_box_size(poly.exterior.coords[:-1])
            if abs(s - length) < abs(b - length):
                length, width = s, b
            else:
                length, width = b, s
        angle = direction2angle(direction) - 90
        angle = angle if angle >= 0 else angle + 360
        height_2_floor = 90 if category == 'window' else 0
        data = {'type': category, 'pos': (round(pos[0], 2), round(pos[1], 2), height_2_floor), 'length': round(length, 2),
                'width': round(width, 2), 'rotate': angle}
        if category in ['door', 'window']:
            height = 150 if category == 'window' else 210
            data['height'] = height
            pretty_res['windowsDoors'].append(data)
        else:
            pretty_res['furniture'].append(data)


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


def plot_layout(polys, labels, save_file=None):
    fig, ax = plt.subplots(figsize=(26, 14), dpi=72, nrows=1, ncols=1)
    plot_polys(ax, polys, labels)
    if save_file:
        plt.savefig(save_file)
    else:
        plt.show()


def calculate_angle_cosine(p1, p2, p3):
    p1 = np.array(p1)
    p2 = np.array(p2)
    p3 = np.array(p3)
    v1 = p2 - p1
    v2 = p2 - p3
    dot_product = np.dot(v1, v2)
    norm_product = np.linalg.norm(v1) * np.linalg.norm(v2)
    if norm_product == 0:
        return 1
    angle_cosine = dot_product / norm_product
    return angle_cosine


def agg_by_val(li):
    li.sort()
    idx = 0
    win_size = 1
    res = []
    while (idx + win_size) <= len(li):
        if (idx + win_size) < len(li) and (li[idx + win_size] - li[idx]) <= win_size:
            win_size += 1
        else:
            res.append(li[idx: idx + win_size])
            idx += win_size
            win_size = 1
    return res


def point_distance(p1, p2):
    return math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)


def get_poly_lines(poly):
    lines = []
    if isinstance(poly, tuple):
        poly = poly.poly
    coords = poly.exterior.coords
    if poly.exterior.is_ccw:
        coords = coords[::-1]
    for i in range(len(coords) - 1):
        lines.append(LineString([coords[i], coords[i + 1]]))
    return lines


def remove_short_edges(points, thres=5.0):
    short_edges = []
    for i in range(len(points)):
        p1 = points[i]
        p2 = points[(i + 1) % len(points)]
        if point_distance(p1, p2) < thres:
            short_edges.append(i)
    short_edge_agg = agg_by_val(short_edges)
    rm_idx = []
    for short_edge in short_edge_agg:
        if len(short_edge) > 1:
            rm_idx.extend(short_edge[1:])
        else:
            i = short_edge[0]
            p1 = points[(i - 1) % len(points)]
            p2 = points[i]
            p3 = points[(i + 1) % len(points)]
            p4 = points[(i + 2) % len(points)]
            if 1 - abs(calculate_angle_cosine(p1, p2, p4)) < 0.1 \
                    or 1 - abs(calculate_angle_cosine(p2, p3, p4)) < 0.1:
                rm_idx.extend([i, (i + 1) % len(points)])
            else:
                new_pos = calculate_cross_pos(p1, p2, p3, p4)
                if point_distance(new_pos, p2) < 2 * thres:
                    points[i] = new_pos
                    rm_idx.append((i + 1) % len(points))
    if len(rm_idx) < 1:
        return points
    for i in sorted(rm_idx, reverse=True):
        del points[i]
    return remove_short_edges(points, thres=thres)


def calculate_cross_pos(p1, p2, p3, p4):
    x1, y1 = p1
    x2, y2 = p2
    x3, y3 = p3
    x4, y4 = p4
    x0 = ((x3 - x4) * (x2 * y1 - x1 * y2) - (x1 - x2) * (x4 * y3 - x3 * y4)) / ((x3 - x4) * (y1 - y2) - (x1 - x2) * (y3 - y4))
    y0 = ((y3 - y4) * (y2 * x1 - y1 * x2) - (y1 - y2) * (y4 * x3 - y3 * x4)) / ((y3 - y4) * (x1 - x2) - (y1 - y2) * (x3 - x4))
    return x0, y0


def post_process_box_data(data):
    res_li = []
    for box in data.obb:
        points = box.xyxyxyxyn[0].tolist()
        lab = data.names[box.cls.cpu().numpy()[0].astype(int)]
        res_li.append({'points': points, 'label': lab, 'prob': box.conf.cpu().numpy()[0].astype(float)})
    return res_li


def post_process_seg_data(data):
    res_li = []
    for mask, box in zip(data.masks, data.boxes):
        if len(mask.xyn) > 1:
            print('multi mask', box.cls)
        points = mask.xyn[0].tolist()
        lab = data.names[box.cls.cpu().numpy()[0].astype(int)]
        res_li.append({'points': points, 'label': lab, 'prob': box.conf.cpu().numpy()[0].astype(float)})
    return res_li


def get_corner_colors(image_array):
    height, width, _ = image_array.shape

    corners = [
        image_array[0, 0],
        image_array[0, width - 1],
        image_array[height - 1, 0],
        image_array[height - 1, width - 1]
    ]

    return corners


def replace_color(image, thresh=10):
    image_array = np.array(image)

    corners = get_corner_colors(image_array)
    most_common_color = np.mean(corners, axis=0)
    white_color = [255, 255, 255]
    mask = np.sqrt(np.sum((image_array - most_common_color) ** 2, axis=-1)) < thresh
    image_array[mask] = white_color
    return Image.fromarray(image_array)


if __name__ == '__main__':
    filename = 'D:\\data\\run\\175a7ed0c9ce48d4b1d9ce446f6d1b0c_yolo.json'
    with open(filename, 'r') as f:
        data = json.load(f)
    print(data)
    pretty_res = post_process_frame_data(data, debug=True)
    with open('simple_layout.json', 'w', encoding='utf8') as f:
        json.dump(pretty_res, f, ensure_ascii=False)
