import numpy as np

from ..environment import torch, cv2, plt




'''
1.visualize_grid_cells(image_path, W, H, X, Y)
将图像划分的grid cell进行可视化
image_path：图像的路径
W:设置图像的宽
H：设置图像的高
X、Y：设置grid cell的数量为X行×Y列个
'''
def visualize_grid_cells(image_path, W, H, ROW, COLUMN, Title='Resized Image with Grid Cells'):
    # 读取图像
    image = cv2.imread(image_path)

    # 如果图像为空（例如文件不存在或损坏），则退出
    if image is None:
        print(f"Error: Could not read image at {image_path}")
        return

        # 调整图像大小
    resized_image = cv2.resize(image, (W, H))

    # 划分grid cell
    cell_width = W // COLUMN
    cell_height = H // ROW

    # 遍历每个grid cell，绘制边界
    for i in range(COLUMN):
        for j in range(ROW):
            # 计算grid cell的左上角坐标
            start_x = i * cell_width
            start_y = j * cell_height
            # 计算grid cell的右下角坐标
            end_x = min((i + 1) * cell_width, W)  # 防止超出图像宽度
            end_y = min((j + 1) * cell_height, H)  # 防止超出图像高度

            # 绘制grid cell边界（使用BGR颜色，因为在OpenCV中图像是BGR格式的）
            cv2.rectangle(resized_image, (start_x, start_y), (end_x, end_y), (0, 255, 0), 2)

            # 显示原图与grid cell（合并到一张图像）
    plt.imshow(cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB))
    plt.title(Title)
    plt.axis('off')
    plt.show()





'''
2.show_images
展示图像
'''
def show_images(imgs, num_rows, num_cols, titles=None, scale=1.5, normalize = False):
    # 绘制图像列表
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    axes = axes.flatten()
    for i, (ax, img) in enumerate(zip(axes, imgs)):
        if torch.is_tensor(img):
            # 图像张量
            if img.dim() == 3:
                img = img.permute(1, 2, 0)  # 将通道维度移动到最后的位置
            if normalize:
                img = img.numpy().clip(0, 1)  # 归一化到0-1范围，并裁剪超出范围的值
            ax.imshow(img)

        else:
            # PIL图像
            ax.imshow(img)
        ax.axes.get_xaxis().set_visible(False)
        ax.axes.get_yaxis().set_visible(False)
        if titles:
            ax.set_title(titles[i])
    return axes





'''
3.box_corner_to_center(boxes)
    summary: it transform the box corners to center
'''
def box_corner_to_center(boxes):
    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    cx = (x1 + x2) // 2
    cy = (y1 + y2) // 2
    w = x2 - x1
    h = y2 - y1
    boxes = torch.stack([cx, cy, w, h], axis=-1)
    return boxes





'''
4.bbox_to_rect
    summary: it draw bounding box rectangle
'''
def bbox_to_rect(bbox, color):
    return plt.Rectangle(xy=(bbox[0], bbox[1]),width=bbox[2]-bbox[0],height=bbox[3]-bbox[1],fill=False,edgecolor=color,linewidth=2)


'''
5.show_boxes(axes, bboxes, labels=None, colors=None)
    summary: it show the box corners
    
'''
def show_boxes(axes, bboxes, labels, colors=None):
    """显示所有的边框"""
    def _make_list(obj, default_values=None):
        if obj is None:
            obj = default_values
        elif not isinstance(obj, (list, tuple)):
            obj = [obj]
        return obj

    labels = _make_list(labels)
    colors = _make_list(colors,['b', 'g', 'r', 'm', 'c' ])
    for i, bbox in enumerate(bboxes):
        color = colors[i % len(colors)]
        rect = bbox_to_rect(bbox.detach().numpy(), color)
        axes.add_patch(rect)
        if labels and len(labels)>i:
            text_color = 'k' if color== 'w' else 'w'
            axes.text(rect.xy[0],rect.xy[1],labels[i],va='center',ha='center',color=text_color,fontsize=9,bbox=dict(facecolor=color, lw=0))


# 5.visualize_bounding_boxes
def visualize_bounding_boxes(image, bounding_boxes, labels):
    """
    Visualizes bounding boxes on an image.

    Parameters:
        image (numpy.ndarray): The input image.
        bounding_boxes (list): A list of bounding boxes in the format [xmin, ymin, xmax, ymax].
        labels (list): A list of labels corresponding to each bounding box.
    """
    # Create figure and axes
    fig, ax = plt.subplots()

    # Display the image
    ax.imshow(image)

    # Add bounding boxes
    for bbox, label in zip(bounding_boxes, labels):
        xmin, ymin, xmax, ymax = bbox
        # Convert relative coordinates to absolute coordinates
        xmin_abs = xmin * image.shape[1]
        ymin_abs = ymin * image.shape[0]
        xmax_abs = xmax * image.shape[1]
        ymax_abs = ymax * image.shape[0]
        width = xmax_abs - xmin_abs
        height = ymax_abs - ymin_abs
        # Create a Rectangle patch
        rect = plt.Rectangle((xmin_abs, ymin_abs), width, height, fill=False, edgecolor='red', linewidth=2)
        ax.add_patch(rect)
        # Add label
        ax.text(xmin_abs, ymin_abs, label, color='red', fontsize=12, verticalalignment='top')

    plt.show()



'''
6.box_iou_xyxy()
    summary: 计算预测框与真实框之间的IOU值，坐标格式为xyxy
    Args:
        boxes (list): A list of bounding boxes in the format [xmin, ymin, xmax, ymax].
'''
def box_iou_xyxy(boxe1, boxe2):
    # 获取box1的具体坐标
    x1min, y1min, x2max, y2max = boxe1[0], boxe1[1], boxe1[2], boxe1[3]

    # 计算boxe1的框面积
    S1 = np.maximum((x2max - x1min) * (y2max - y1min), 0)

    # 获取box2的具体坐标
    x2min, y2min, x2max, y2max = boxe2[0], boxe2[1], boxe2[2], boxe2[3]

    # 计算boxe2的框面积.
    S2 = np.maximum((x2max - x2min) * (y2max - y2min), 0)

    # 计算相交矩形框的坐标，即boxe1与boxe2的交集
    xmin = np.maximum(x1min, x2min)
    ymin = np.maximum(y1min, y2min)
    xmax = np.minimum(x2max, x2max)
    ymax = np.minimum(y2max, y2max)

    # 计算两个矩形框的交集
    intersection = np.maximum((xmax - xmin) * (ymax - ymin), 0)

    IoU = intersection / np.maximum((S1 + S2 - intersection),1.)

    return IoU



'''
7.box_iou_xywh()
'''
def box_iou_xywh(boxe1, boxe2):
    # 将xywh化为xyxy
    x1min, y1min, x1max, y1max = boxe1[0] - boxe1[2]/2.0 , boxe1[1] - boxe1[3]/2.0, boxe1[0] + boxe1[2]/2.0 , boxe1[1] + boxe1[3]/2.0
    x2min, y2min, x2max, y2max = boxe2[0] - boxe2[2]/2.0 , boxe2[1] - boxe2[3]/2.0, boxe2[0] + boxe2[2]/2.0 , boxe2[1] + boxe2[3]/2.0

    return box_iou_xyxy((x1min, y1min, x1max, y1max), (x2min, y2min, x2max, y2max))




'''
8.draw_anchor_boxes
'''
def draw_anchor_boxes(image, center_x, center_y, widths, heights):
    """
    在图像中画出给定中心点和多种宽高的锚框。

    参数：
        image: 输入的图像，可以是 numpy 数组。
        center_x: 中心点的 x 坐标。
        center_y: 中心点的 y 坐标。
        widths: 锚框的宽度列表。
        heights: 锚框的高度列表。

    返回值：
        annotated_image: 画有锚框的图像，numpy 数组格式。
    """
    # 创建一个副本以保留原始图像
    annotated_image = image.copy()

    # 循环绘制每个锚框
    for width, height in zip(widths, heights):
        # 计算左上角和右下角坐标
        top_left = (int(center_x - width / 2), int(center_y - height / 2))
        bottom_right = (int(center_x + width / 2), int(center_y + height / 2))

        # 在图像上画出锚框
        annotated_image = cv2.rectangle(annotated_image, top_left, bottom_right, (0, 255, 0), 2)

    return annotated_image









'''
9.get_objectness_label(img, gt_boxes, gtlabels, iou_threshold=0.7, anchors = [])
    summary: 标注预测框的objectness
    img 是输入的图像数据，形状是[N, C, H, W]
    gt_boxes，真实框，维度是[N, 50, 4]，其中50是真实框数目的上限，当图片中真实框不足50个时，不足部分的坐标全为0
              真实框坐标格式是xywh，这里使用相对值
    gt_labels，真实框所属类别，维度是[N, 50]
    iou_threshold，当预测框与真实框的iou大于iou_threshold时不将其看作是负样本
    anchors，锚框可选的尺寸
    anchor_masks，通过与anchors一起确定本层级的特征图应该选用多大尺寸的锚框
    num_classes，类别数目
    downsample，特征图相对于输入网络的图片尺寸变化的比例
'''
def get_objectness_label(img, gt_boxes, gtlabels, iou_threshold=0.7, anchors = [116, 90, 156, 198, 373, 326], num_classes = 3, downsample = 32):
    img_shape = img.shape
    batch_size = img.shape[0]
    num_anchors = len(anchors) // 2
    input_h = img.shape[2]
    input_w = img.shape[3]

    # 将输入图片划分成num_rows x num_cols个候选区域，每个候选区域的边长是downsample
    # 计算共有多少列候选区域
    num_cols = input_w // downsample
    # 计算共有多少行候选区域
    num_rows = input_h // downsample

    label_objectness = np.zeros([batch_size, num_anchors, num_rows, num_cols])
    label_classification = np.zeros([batch_size, num_anchors, num_classes, num_rows, num_cols])
    label_location = np.zeros([batch_size, num_anchors, 4, num_rows, num_cols])

    # 这里有个scale_location，用来调节大目标和小目标的权重的
    # 大目标和小目标对于相同的偏移，大目标更加敏感，故增强这个权重系数
    scale_location = np.zeros([batch_size, num_anchors, num_rows, num_cols])

    # batch_size进行循环，处理每一张图片
    for n in range(batch_size):
        # 对图片上的真实框进行遍历，依次找出与真实框形状最匹配的锚框
        for n_gt in range(len(gt_boxes[n])):
            gt_box = gt_boxes[n][n_gt]
            # 真实框所属类别
            gt_label = gtlabels[n][n_gt]
            # 真实框的中心坐标和长宽（YOLO中采用xywh）
            gt_center_x = gt_box[0]
            gt_center_y = gt_box[1]
            gt_width = gt_box[2]
            gt_height = gt_box[3]
            if (gt_width < 1e-3) or (gt_height < 1e-3):
                continue

            # 计算这个真实框的中心点落在哪个格子里,xy是真实边界框的中心点在图像中的归一化坐标（即范围在 [0, 1] 内）
            i = int(gt_center_y * num_rows)
            j = int(gt_center_x * num_cols)

            ious = []

            for ka in range(num_anchors):
                # x，y为0，假设锚框和真实框的x，y相同
                gt_bbox = [0., 0., float(gt_width), float(gt_height)]
                # 生成输入（i，j）这个格子里的锚框
                anchor_w = anchors[ka * 2]
                anchor_h = anchors[ka * 2 + 1]
                anchor_bbox = [0., 0., anchor_w / float(input_w), anchor_h / float(input_h)]
                # 计算iou
                iou = box_iou_xywh(gt_bbox, anchor_bbox)
                ious.append(iou)
            ious = np.array(ious)
            inds = np.argsort(ious)
            k = inds[-1]
            label_objectness[n, k, i, j] = 1
            c = gt_label
            # 使用one-hot方法标注
            label_classification[n, k, c, i, j] = 1

            # for those prediction bbox with objectness =1, set label of location
            # 反算出真实框的tx*，ty*，tw*，th*
            dx_label = gt_center_x * num_cols - j
            dy_label = gt_center_y * num_rows - i
            dw_label = np.log(gt_width * input_w / anchors[k * 2])
            dh_label = np.log(gt_height * input_h / anchors[k * 2 + 1])
            label_location[n, k, 0, i, j] = dx_label
            label_location[n, k, 1, i, j] = dy_label
            label_location[n, k, 2, i, j] = dw_label
            label_location[n, k, 3, i, j] = dh_label

            # scale_location用来调节不同尺寸的锚框对损失函数的贡献，作为加权系数和位置损失函数相乘
            # 根据这个计算方法，可以得出真实框越大，权重越小
            scale_location[n, k, i, j] = 2.0 - gt_width * gt_height

    # 目前根据每张图片上所有出现过的gt_box都标注出了objectness为正的预测框，剩下的预测框的objectness为0
    # 对于objectness为1的预测框，标出它们所包含的物体类别以及位置回归的目标
    return label_objectness.astype('float32'), label_location.astype('float32'), label_classification.astype('float32'), scale_location.astype('float32')






'''
10.NMS非极大值抑制
    summary: Non-Maximum Suppression
'''
def nms(bboxes, scores, score_thresh, nms_thresh, pre_nms_topk, i=0, c=0):
    inds = np.argsort(scores)
    inds = inds[::-1]
    keep_inds = []
    while(len(inds) > 0):
        cur_ind = inds[0]
        cur_score = scores[cur_ind]
        # 如果该框的分数小于score.thresh，则丢弃该框
        if cur_score < score_thresh:
            break

        keep = True
        for ind in keep_inds:
            current_box = bboxes[cur_ind]
            remain_box = bboxes[ind]
            iou = box_iou_xywh(current_box, remain_box)
            if iou > nms_thresh:
                keep = False
                break
        # 调试代码，用于观察信息使用
        if i == 0 and c == 4 and cur_ind == 951:
            print('suppressed, ',keep,i,c,cur_ind,ind,iou)
        if keep:
            keep_inds.append(cur_ind)
        # 去除当前处理的边界框索引，即将 inds 数组中的第一个元素移除，继续处理下一个边界框。
        inds = inds[1:]
    return np.array(keep_inds)






'''
11.multiclass_nms
    summary: multiclass-Non-Maximum Suppression
'''
def multiclass_nms(bboxes, scores, score_thresh = 0.01, nms_thresh = 0.45, pre_nms_topk = 1000, pos_nms_topk = 100 ):
    # 获取批量大小和类别数量
    batch_size = bboxes.shape[0]
    class_num = scores.shape[1]
    rets = []
    # 遍历每个样本
    for i in range(batch_size):
        # 获取当前样本的边界框和分数
        bboxes_i = bboxes[i]
        scores_i = scores[i]
        ret = []
        # 遍历每个类别
        for c in range(class_num):
            # 获取当前类别的分数
            scores_i_c = scores_i[c]
            # 使用 NMS 函数进行非最大抑制
            keep_inds = nms(bboxes_i, scores_i_c, score_thresh, nms_thresh, pre_nms_topk, i, c)

            # 如果保留的边界框数量小于1，则跳过当前类别
            if len(keep_inds) < 1:
                continue
            # 根据保留的索引获取对应的边界框和分数
            keep_bboxes = bboxes_i[keep_inds]
            keep_scores = scores_i_c[keep_inds]
            # 构造保留结果
            keep_results = np.zeros([keep_scores.shape[0], 6])
            keep_results[:, 0] = c # 类别索引
            keep_results[:, 1] = keep_scores[:] # 分数
            keep_results[:, 2:6] = keep_bboxes[:, :]    # 边界框坐标
            ret.append(keep_results)
        # 如果当前样本没有保留的结果，则跳过
        if len(ret) < 1:
            rets.append(ret)
            continue
        # 将每个类别的结果拼接在一起
        ret_i = np.concatenate(ret, axis=0)
        # 如果保留的结果数量超过了指定的阈值，则根据分数进行排序并截取指定数量的结果
        scores_i = ret_i[:, 1]
        if len(scores_i) > pos_nms_topk:
            inds = np.argsort(scores_i)[::-1]
            inds = inds[:pos_nms_topk]
            ret_i = ret_i[inds]

        rets.append(ret_i)

    return rets


























































