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


def c(score):
    # 根据分数返回颜色的函数，这里仅作示例，您可以根据需要修改
    return (1, 0, 0) if score > 0.9 else (0, 1, 0)


def postprocess(lines, scores, diag_threshold, min_score, remove_overlaps):
    # 假设的后处理函数，用于过滤线段
    nlines = []
    nscores = []
    for line, score in zip(lines, scores):
        if score >= min_score:
            nlines.append(line)
            nscores.append(score)
    return np.array(nlines), np.array(nscores)


def show_line(img, pred, epoch, writer):
    im = img.permute(1, 2, 0).cpu().numpy()

    # 绘制边界框
    boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), pred[0]["boxes"],
                                      colors="yellow", width=1).permute(1, 2, 0).cpu().numpy()

    H = pred[-1]['wires']
    lines = H["lines"][0].cpu().numpy() / 128 * im.shape[:2]
    scores = H["score"][0].cpu().numpy()

    print(f"Lines before deduplication: {len(lines)}")

    # 移除重复的线段
    for i in range(1, len(lines)):
        if (lines[i] == lines[0]).all():
            lines = lines[:i]
            scores = scores[:i]
            break

    print(f"Lines after deduplication: {len(lines)}")

    # 后处理线段以移除重叠的线段
    diag = (im.shape[0] ** 2 + im.shape[1] ** 2) ** 0.5
    nlines, nscores = postprocess(lines, scores, diag * 0.01, 0, False)

    print(f"Lines after postprocessing: {len(nlines)}")

    # 创建一个新的图像并绘制线段和边界框
    fig, ax = plt.subplots(figsize=(boxed_image.shape[1] / 100, boxed_image.shape[0] / 100))
    ax.imshow(boxed_image)
    ax.set_axis_off()
    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
    plt.margins(0, 0)
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    plt.gca().yaxis.set_major_locator(plt.NullLocator())

    PLTOPTS = {"color": "#33FFFF", "s": 15, "edgecolors": "none", "zorder": 5}
    for (a, b), s in zip(nlines, nscores):
        if s < 0.85:  # 调整阈值以筛选显示的线段
            continue
        plt.plot([a[1], b[1]], [a[0], b[0]], c=c(s), linewidth=2, zorder=s)
        plt.scatter(a[1], a[0], **PLTOPTS)
        plt.scatter(b[1], b[0], **PLTOPTS)

    plt.tight_layout()
    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("output_with_boxes_and_lines", img2, epoch)
    print("Image with boxes and lines added to TensorBoard.")