import time

import torch
import matplotlib.pyplot as plt
import numpy as np
from torchvision import transforms

from models.wirenet.postprocess import postprocess


def box_line(pred):
    '''
    :param pred: 预测结果
    :return:

    box与line一一对应
{'box': [0.0, 34.23157501220703, 151.70858764648438, 125.10173797607422], 'line': array([[ 1.9720564, 81.73457  ],
[ 1.9933801, 41.730167 ]], dtype=float32)}
    '''
    box_line = [[] for _ in range((len(pred) - 1))]
    for idx, box_ in enumerate(pred[0:-1]):
        box = box_['boxes']  # 是一个tensor
        line = pred[-1]['wires']['lines'][idx].cpu().numpy() / 128 * 512
        score = pred[-1]['wires']['score'][idx]
        for i in box:
            aaa = {}
            aaa['box'] = i.tolist()
            aaa['line'] = []
            score_max = 0.0
            for j in range(len(line)):
                if (line[j][0][0] >= i[0] and line[j][1][0] >= i[0] and line[j][0][0] <= i[2] and
                        line[j][1][0] <= i[2] and line[j][0][1] >= i[1] and line[j][1][1] >= i[1] and
                        line[j][0][1] <= i[3] and line[j][1][1] <= i[3]):
                    if score[j] > score_max:
                        aaa['line'] = line[j]
                        score_max = score[j]
            box_line[idx].append(aaa)


def box_line_(pred):
    for idx, box_ in enumerate(pred[0:-1]):
        box = box_['boxes']  # 是一个tensor
        line = pred[-1]['wires']['lines'][idx].cpu().numpy() / 128 * 512
        score = pred[-1]['wires']['score'][idx]
        line_ = []
        score_ = []

        for i in box:
            score_max = 0.0
            tmp = [[0.0, 0.0], [0.0, 0.0]]

            for j in range(len(line)):
                if (line[j][0][1] >= i[0] and line[j][1][1] >= i[0] and
                        line[j][0][1] <= i[2] and line[j][1][1] <= i[2] and
                        line[j][0][0] >= i[1] and line[j][1][0] >= i[1] and
                        line[j][0][0] <= i[3] and line[j][1][0] <= i[3]):

                    if score[j] > score_max:
                        tmp = line[j]
                        score_max = score[j]
            line_.append(tmp)
            score_.append(score_max)
        processed_list = torch.tensor(line_)
        pred[idx]['line'] = processed_list

        processed_s_list = torch.tensor(score_)
        pred[idx]['line_score'] = processed_s_list
    return pred


# box与line匹配后画在一张图上，不设置阈值，直接画
def show_(imgs, pred, epoch, writer):
    col = [
        '#1f77b4', '#aec7e8', '#ff7f0e', '#ffbb78', '#2ca02c',
        '#98df8a', '#d62728', '#ff9896', '#9467bd', '#c5b0d5',
        '#8c564b', '#c49c94', '#e377c2', '#f7b6d2', '#7f7f7f',
        '#c7c7c7', '#bcbd22', '#dbdb8d', '#17becf', '#9edae5',
        '#8dd3c7', '#ffffb3', '#bebada', '#fb8072', '#80b1d3',
        '#fdb462', '#b3de69', '#fccde5', '#bc80bd', '#ccebc5',
        '#ffed6f', '#8da0cb', '#e78ac3', '#e5c494', '#b3b3b3',
        '#fdbf6f', '#ff7f00', '#cab2d6', '#637939', '#b5cf6b',
        '#cedb9c', '#8c6d31', '#e7969c', '#d6616b', '#7b4173',
        '#ad494a', '#843c39', '#dd8452', '#f7f7f7', '#cccccc',
        '#969696', '#525252', '#f7fcfd', '#e5f5f9', '#ccece6',
        '#99d8c9', '#66c2a4', '#2ca25f', '#008d4c', '#005a32',
        '#f7fcf0', '#e0f3db', '#ccebc5', '#a8ddb5', '#7bccc4',
        '#4eb3d3', '#2b8cbe', '#08589e', '#f7fcfd', '#e0ecf4',
        '#bfd3e6', '#9ebcda', '#8c96c6', '#8c6bb4', '#88419d',
        '#810f7c', '#4d004b', '#f7f7f7', '#efefef', '#d9d9d9',
        '#bfbfbf', '#969696', '#737373', '#525252', '#252525',
        '#000000', '#ffffff', '#ffeda0', '#fed976', '#feb24c',
        '#fd8d3c', '#fc4e2a', '#e31a1c', '#bd0026', '#800026',
        '#ff6f61', '#ff9e64', '#ff6347', '#ffa07a', '#fa8072'
    ]
    # print(len(col))
    im = imgs[0].permute(1, 2, 0)
    boxes = pred[0]['boxes'].cpu().numpy()
    line = pred[0]['line'].cpu().numpy()

    # 可视化预测结
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(np.array(im))

    for idx, box in enumerate(boxes):
        x0, y0, x1, y1 = box
        ax.add_patch(
            plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=col[idx], linewidth=1))

    for idx, (a, b) in enumerate(line):
        ax.scatter(a[1], a[0], c='#871F78', s=2)
        ax.scatter(b[1], b[0], c='#871F78', s=2)
        ax.plot([a[1], b[1]], [a[0], b[0]], c=col[idx], linewidth=1)

    # 将Matplotlib图像转换为Tensor
    fig.canvas.draw()
    image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8).reshape(
        fig.canvas.get_width_height()[::-1] + (3,))
    plt.close()
    img2 = transforms.ToTensor()(image_from_plot)

    writer.add_image("all", img2, epoch)


# box与line匹配后画在一张图上，设置阈值
def show_predict(imgs, pred, t_start):
    col = [
        '#1f77b4', '#aec7e8', '#ff7f0e', '#ffbb78', '#2ca02c',
        '#98df8a', '#d62728', '#ff9896', '#9467bd', '#c5b0d5',
        '#8c564b', '#c49c94', '#e377c2', '#f7b6d2', '#7f7f7f',
        '#c7c7c7', '#bcbd22', '#dbdb8d', '#17becf', '#9edae5',
        '#8dd3c7', '#ffffb3', '#bebada', '#fb8072', '#80b1d3',
        '#fdb462', '#b3de69', '#fccde5', '#bc80bd', '#ccebc5',
        '#ffed6f', '#8da0cb', '#e78ac3', '#e5c494', '#b3b3b3',
        '#fdbf6f', '#ff7f00', '#cab2d6', '#637939', '#b5cf6b',
        '#cedb9c', '#8c6d31', '#e7969c', '#d6616b', '#7b4173',
        '#ad494a', '#843c39', '#dd8452', '#f7f7f7', '#cccccc',
        '#969696', '#525252', '#f7fcfd', '#e5f5f9', '#ccece6',
        '#99d8c9', '#66c2a4', '#2ca25f', '#008d4c', '#005a32',
        '#f7fcf0', '#e0f3db', '#ccebc5', '#a8ddb5', '#7bccc4',
        '#4eb3d3', '#2b8cbe', '#08589e', '#f7fcfd', '#e0ecf4',
        '#bfd3e6', '#9ebcda', '#8c96c6', '#8c6bb4', '#88419d',
        '#810f7c', '#4d004b', '#f7f7f7', '#efefef', '#d9d9d9',
        '#bfbfbf', '#969696', '#737373', '#525252', '#252525',
        '#000000', '#ffffff', '#ffeda0', '#fed976', '#feb24c',
        '#fd8d3c', '#fc4e2a', '#e31a1c', '#bd0026', '#800026',
        '#ff6f61', '#ff9e64', '#ff6347', '#ffa07a', '#fa8072'
    ]
    print(len(col))
    im = imgs.permute(1, 2, 0)  # 处理为 [512, 512, 3]
    boxes = pred[0]['boxes'].cpu().numpy()
    box_scores = pred[0]['scores'].cpu().numpy()
    lines = pred[0]['line'].cpu().numpy()
    line_scores = pred[0]['line_score'].cpu().numpy()

    # 可视化预测结
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(np.array(im))
    idx = 0

    for box, line, box_score, line_score in zip(boxes, lines, box_scores, line_scores):
        x0, y0, x1, y1 = box
        a, b = line
        if box_score > 0.7 and line_score > 0.9:
            ax.add_patch(
                plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=col[idx], linewidth=1))
            ax.scatter(a[1], a[0], c='#871F78', s=10)
            ax.scatter(b[1], b[0], c='#871F78', s=10)
            ax.plot([a[1], b[1]], [a[0], b[0]], c=col[idx], linewidth=1)
            idx = idx + 1
    t_end = time.time()
    print(f'predict used:{t_end - t_start}')

    plt.show()


# 下面的都没有进行box与line的一一匹配
# 只画线，设阈值
def show_line(imgs, pred, t_start):

    im = imgs.permute(1, 2, 0)
    line = pred[-1]['wires']['lines'][0].cpu().numpy() / 128 * 512
    # print(pred[-1]['wires']['score'])
    line_score = pred[-1]['wires']['score'].cpu().numpy()[0]

    # 可视化预测结
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(np.array(im))

    for idx, (a, b) in enumerate(line):
        if line_score[idx] < 0.9:
            continue
        ax.scatter(a[1], a[0], c='#871F78', s=2)
        ax.scatter(b[1], b[0], c='#871F78', s=2)
        ax.plot([a[1], b[1]], [a[0], b[0]], c='red', linewidth=1)

    t_end = time.time()
    print(f'show_line used:{t_end - t_start}')

    plt.show()


# 只画框，设阈值
def show_box(imgs, pred, t_start):
    col = [
        '#1f77b4', '#aec7e8', '#ff7f0e', '#ffbb78', '#2ca02c',
        '#98df8a', '#d62728', '#ff9896', '#9467bd', '#c5b0d5',
        '#8c564b', '#c49c94', '#e377c2', '#f7b6d2', '#7f7f7f',
        '#c7c7c7', '#bcbd22', '#dbdb8d', '#17becf', '#9edae5',
        '#8dd3c7', '#ffffb3', '#bebada', '#fb8072', '#80b1d3',
        '#fdb462', '#b3de69', '#fccde5', '#bc80bd', '#ccebc5',
        '#ffed6f', '#8da0cb', '#e78ac3', '#e5c494', '#b3b3b3',
        '#fdbf6f', '#ff7f00', '#cab2d6', '#637939', '#b5cf6b',
        '#cedb9c', '#8c6d31', '#e7969c', '#d6616b', '#7b4173',
        '#ad494a', '#843c39', '#dd8452', '#f7f7f7', '#cccccc',
        '#969696', '#525252', '#f7fcfd', '#e5f5f9', '#ccece6',
        '#99d8c9', '#66c2a4', '#2ca25f', '#008d4c', '#005a32',
        '#f7fcf0', '#e0f3db', '#ccebc5', '#a8ddb5', '#7bccc4',
        '#4eb3d3', '#2b8cbe', '#08589e', '#f7fcfd', '#e0ecf4',
        '#bfd3e6', '#9ebcda', '#8c96c6', '#8c6bb4', '#88419d',
        '#810f7c', '#4d004b', '#f7f7f7', '#efefef', '#d9d9d9',
        '#bfbfbf', '#969696', '#737373', '#525252', '#252525',
        '#000000', '#ffffff', '#ffeda0', '#fed976', '#feb24c',
        '#fd8d3c', '#fc4e2a', '#e31a1c', '#bd0026', '#800026',
        '#ff6f61', '#ff9e64', '#ff6347', '#ffa07a', '#fa8072'
    ]
    # print(len(col))
    im = imgs.permute(1, 2, 0)
    boxes = pred[0]['boxes'].cpu().numpy()
    box_scores = pred[0]['scores'].cpu().numpy()

    # 可视化预测结
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(np.array(im))

    for idx, box in enumerate(boxes):
        if box_scores[idx] < 0.7:
            continue
        x0, y0, x1, y1 = box
        ax.add_patch(
            plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=col[idx], linewidth=1))

    t_end = time.time()
    print(f'show_box used:{t_end - t_start}')

    plt.show()


# 将show_line与show_box合并，传入参数确定显示框还是线  都不显示，输出原图
def show_box_or_line(imgs, pred, show_line=False, show_box=False):
    col = [
        '#1f77b4', '#aec7e8', '#ff7f0e', '#ffbb78', '#2ca02c',
        '#98df8a', '#d62728', '#ff9896', '#9467bd', '#c5b0d5',
        '#8c564b', '#c49c94', '#e377c2', '#f7b6d2', '#7f7f7f',
        '#c7c7c7', '#bcbd22', '#dbdb8d', '#17becf', '#9edae5',
        '#8dd3c7', '#ffffb3', '#bebada', '#fb8072', '#80b1d3',
        '#fdb462', '#b3de69', '#fccde5', '#bc80bd', '#ccebc5',
        '#ffed6f', '#8da0cb', '#e78ac3', '#e5c494', '#b3b3b3',
        '#fdbf6f', '#ff7f00', '#cab2d6', '#637939', '#b5cf6b',
        '#cedb9c', '#8c6d31', '#e7969c', '#d6616b', '#7b4173',
        '#ad494a', '#843c39', '#dd8452', '#f7f7f7', '#cccccc',
        '#969696', '#525252', '#f7fcfd', '#e5f5f9', '#ccece6',
        '#99d8c9', '#66c2a4', '#2ca25f', '#008d4c', '#005a32',
        '#f7fcf0', '#e0f3db', '#ccebc5', '#a8ddb5', '#7bccc4',
        '#4eb3d3', '#2b8cbe', '#08589e', '#f7fcfd', '#e0ecf4',
        '#bfd3e6', '#9ebcda', '#8c96c6', '#8c6bb4', '#88419d',
        '#810f7c', '#4d004b', '#f7f7f7', '#efefef', '#d9d9d9',
        '#bfbfbf', '#969696', '#737373', '#525252', '#252525',
        '#000000', '#ffffff', '#ffeda0', '#fed976', '#feb24c',
        '#fd8d3c', '#fc4e2a', '#e31a1c', '#bd0026', '#800026',
        '#ff6f61', '#ff9e64', '#ff6347', '#ffa07a', '#fa8072'
    ]
    # print(len(col))
    im = imgs.permute(1, 2, 0)
    boxes = pred[0]['boxes'].cpu().numpy()
    line = pred[-1]['wires']['lines'][0].cpu().numpy() / 128 * 512

    # 可视化预测结
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(np.array(im))

    if show_box:
        for idx, box in enumerate(boxes):
            x0, y0, x1, y1 = box
            ax.add_patch(
                plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=col[idx], linewidth=1))

    if show_line:
        for idx, (a, b) in enumerate(line):
            ax.scatter(a[1], a[0], c='#871F78', s=2)
            ax.scatter(b[1], b[0], c='#871F78', s=2)
            ax.plot([a[1], b[1]], [a[0], b[0]], c='red', linewidth=1)

    plt.show()


# 将show_line与show_box合并，传入参数确定显示框还是线  一起画
def show_box_and_line(imgs, pred, show_line=False, show_box=False):
    col = [
        '#1f77b4', '#aec7e8', '#ff7f0e', '#ffbb78', '#2ca02c',
        '#98df8a', '#d62728', '#ff9896', '#9467bd', '#c5b0d5',
        '#8c564b', '#c49c94', '#e377c2', '#f7b6d2', '#7f7f7f',
        '#c7c7c7', '#bcbd22', '#dbdb8d', '#17becf', '#9edae5',
        '#8dd3c7', '#ffffb3', '#bebada', '#fb8072', '#80b1d3',
        '#fdb462', '#b3de69', '#fccde5', '#bc80bd', '#ccebc5',
        '#ffed6f', '#8da0cb', '#e78ac3', '#e5c494', '#b3b3b3',
        '#fdbf6f', '#ff7f00', '#cab2d6', '#637939', '#b5cf6b',
        '#cedb9c', '#8c6d31', '#e7969c', '#d6616b', '#7b4173',
        '#ad494a', '#843c39', '#dd8452', '#f7f7f7', '#cccccc',
        '#969696', '#525252', '#f7fcfd', '#e5f5f9', '#ccece6',
        '#99d8c9', '#66c2a4', '#2ca25f', '#008d4c', '#005a32',
        '#f7fcf0', '#e0f3db', '#ccebc5', '#a8ddb5', '#7bccc4',
        '#4eb3d3', '#2b8cbe', '#08589e', '#f7fcfd', '#e0ecf4',
        '#bfd3e6', '#9ebcda', '#8c96c6', '#8c6bb4', '#88419d',
        '#810f7c', '#4d004b', '#f7f7f7', '#efefef', '#d9d9d9',
        '#bfbfbf', '#969696', '#737373', '#525252', '#252525',
        '#000000', '#ffffff', '#ffeda0', '#fed976', '#feb24c',
        '#fd8d3c', '#fc4e2a', '#e31a1c', '#bd0026', '#800026',
        '#ff6f61', '#ff9e64', '#ff6347', '#ffa07a', '#fa8072'
    ]
    # print(len(col))
    im = imgs.permute(1, 2, 0)
    boxes = pred[0]['boxes'].cpu().numpy()
    line = pred[-1]['wires']['lines'][0].cpu().numpy() / 128 * 512

    # 可视化预测结
    fig, axs = plt.subplots(1, 2, figsize=(10, 10))

    if show_box:
        axs[0].imshow(np.array(im))
        for idx, box in enumerate(boxes):
            x0, y0, x1, y1 = box
            axs[0].add_patch(
                plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=col[idx], linewidth=1))
        axs[0].set_title('Boxes')

    if show_line:
        axs[1].imshow(np.array(im))
        for idx, (a, b) in enumerate(line):
            axs[1].scatter(a[1], a[0], c='#871F78', s=2)
            axs[1].scatter(b[1], b[0], c='#871F78', s=2)
            axs[1].plot([a[1], b[1]], [a[0], b[0]], c='red', linewidth=1)
        axs[1].set_title('Lines')

    # 调整子图之间的距离，防止标题和标签重叠
    plt.tight_layout()
    plt.show()