import math

import numpy as np
import torch


def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int, scale_ratio):

    # 创建yolo_det_clone内容为x1, y1, x2, y2, conf, cls, unlabel (unlabel代表该类是否需要忽略，0：不忽略 其他：忽略)
    yolo_det_clone = yolo_det.copy()
    tmp_0_tensor = np.zeros([len(yolo_det), 1])
    yolo_det_clone = np.concatenate([yolo_det_clone, tmp_0_tensor], axis=1)

    # cls为需要计算的类别
    yolo_det = yolo_det[yolo_det[:, -1] == cls]

    # new_yolo_det为膨胀后数据，内容为x1, y1, x2, y2, flag (flag代表膨胀后车位内是否包含角点 且 与角点方向差值小于90度, 其值为第一个满足条件的角点索引)
    new_yolo_det = np.zeros([len(yolo_det), 7])

    # yolo框膨胀，长的边两边各膨胀0.4倍总长，短的边两边各膨胀0.2倍总长
    x_length = yolo_det[:, 2] - yolo_det[:, 0]      #x2-x1
    y_length = yolo_det[:, 3] - yolo_det[:, 1]     #y2-y1

    # x, y哪个方向差值大哪个方向膨胀的多
    x_dilate_coefficient = ((x_length > y_length) + 1)*scale_ratio
    y_dilate_coefficient = ((~(x_length > y_length)) + 1)*scale_ratio

    # 原始框中心点x_c, y_c
    new_yolo_det[:, 5] = (yolo_det[:, 0] + yolo_det[:, 2]) / 2
    new_yolo_det[:, 6] = (yolo_det[:, 1] + yolo_det[:, 3]) / 2

    # 膨胀
    new_yolo_det[:, 0] = np.round(yolo_det[:, 0] - x_dilate_coefficient * x_length).clip(0, img_shape[1])   #x1 膨胀
    new_yolo_det[:, 1] = np.round(yolo_det[:, 1] - y_dilate_coefficient * y_length).clip(0, img_shape[0])   #y1 膨胀
    new_yolo_det[:, 2] = np.round(yolo_det[:, 2] + x_dilate_coefficient * x_length).clip(0, img_shape[1])   #x2 膨胀
    new_yolo_det[:, 3] = np.round(yolo_det[:, 3] + y_dilate_coefficient * y_length).clip(0, img_shape[0])   #y2 膨胀

    # 判断膨胀后yolo框包含角点关系 && 包含角点的时候计算水平框中心点与角点的角度关系
    # for i in range(0, len(new_yolo_det)):
    #     for j in range(0, len(dmpr_det)):
    #         if new_yolo_det[i, 4] == 0:
    #             [x_p, y_p] = dmpr_det[j, 1:3]
    #             [x1, y1, x2, y2] = new_yolo_det[i, :4]
    #             x_c = (x1+x2)/2
    #             y_c = (y1+y2)/2
    #             if (x_p >= x1) and (x_p <= x2) and (y_p >= y1) and (y_p <= y2):
    #                 direction1 = math.atan2(y_c-y_p, x_c-x_p) / math.pi * 180
    #                 direction2 = dmpr_det[j, 3] / math.pi * 180
    #                 ang_diff = direction1 - direction2
    #                 # direction ∈ （-180， 180） 若角差大于180，需算补角
    #                 if (ang_diff >= -90) and (ang_diff <= 90):
    #                     new_yolo_det[i, 4] = j + 1                  #为防止 j = 0 时赋值，故作 +1 操作
    #                 elif (ang_diff > 180) and (360 - ang_diff <= 90):
    #                     new_yolo_det[i, 4] = j + 1
    #                 elif (ang_diff < -180) and (360 + ang_diff <= 90):
    #                     new_yolo_det[i, 4] = j + 1
    m, n = new_yolo_det.size, dmpr_det.size
    if not m or not n:
        return yolo_det_clone, new_yolo_det

    new_yolo = new_yolo_det[:, np.newaxis, :].repeat(dmpr_det.shape[0], 1)      # 扩展为 (m , n, 5)
    dmpr_det = dmpr_det[np.newaxis, ...].repeat(new_yolo_det.shape[0], 0)
    yolo_dmpr = np.concatenate((new_yolo, dmpr_det), axis=2)              #  (m, n, 10)

    x_p, y_p = yolo_dmpr[..., 8], yolo_dmpr[..., 9]
    x1, y1, x2, y2 = yolo_dmpr[..., 0], yolo_dmpr[..., 1], yolo_dmpr[..., 2], yolo_dmpr[..., 3]
    x_c, y_c = yolo_dmpr[..., 5], yolo_dmpr[..., 6]

    direction1 = np.arctan2(y_c - y_p, x_c - x_p) / math.pi * 180
    direction2 = yolo_dmpr[..., 10] / math.pi * 180
    # direction3 = (direction2 + 90) if (direction2 + 90) <= 180 else (direction2 - 270)
    direction3 = direction2 + 90                    # L形角点另外一个方向
    direction3[direction3 > 180] -= 360
    ang_diff = direction1 - direction2
    ang_diff2 = direction1 - direction3

    # 判断膨胀后yolo框包含角点关系 & & 包含角点的时候计算水平框中心点与角点的角度关系
    # direction ∈ （-180， 180） 若角差大于180，需算补角
    # T形角点比较一个方向，L形角点比较两个方向
    mask = (x_p >= x1) & (x_p <= x2) & (y_p >= y1) & (y_p <= y2) & \
           (((yolo_dmpr[..., 11] <= 0.5) &                          # T形角点情况
                (((ang_diff >= -90) & (ang_diff <= 90)) | ((ang_diff > 180) & ((360 - ang_diff) <= 90)) |
                    (((ang_diff) < -180) & ((360 + ang_diff) <= 90)))) |
            ((yolo_dmpr[..., 11] > 0.5) &                            # L形角点情况
                (((ang_diff >= -90) & (ang_diff <= 90)) | ((ang_diff > 180) & ((360 - ang_diff) <= 90)) |
                    (((ang_diff) < -180) & ((360 + ang_diff) <= 90))) &
                (((ang_diff2 >= -90) & (ang_diff2 <= 90)) | ((ang_diff2 > 180) & ((360 - ang_diff2) <= 90)) |
                    (((ang_diff2) < -180) & ((360 + ang_diff2) <= 90)))))

    res = np.sum(mask, axis=1)

    # 索引两次更新tensor test1
    # yolo_det_clone[yolo_det_clone[:, -2] == cls][:, -1] = new_yolo_det[:, 4]

    # 索引两次更新tensor test2
    # a = [x for x in torch.arange(len(new_yolo_det))]
    # b = [6 for _ in torch.arange(len(new_yolo_det))]
    # index = (torch.LongTensor(a), torch.LongTensor(b))
    # value = torch.tensor(new_yolo_det[:, 4], device=device_)
    # yolo_det_clone[yolo_det_clone[:, -2] == cls].index_put_(index, value)

    yolo_det_clone[yolo_det_clone[:, -2] == cls, -1] = res

    return yolo_det_clone, new_yolo_det