import cv2
import numpy as np
import torch
from loguru import logger

def security_post_process(preds, pars):
    # pars={'solar':0}
    '''
    将光伏板上覆盖物、裂缝识别出来
    '''
    # print(preds[0])
    #logger.info('\n分类结果返回：%s'%preds)
    if not preds[0]:
        return [[],'']
    preds = np.array(preds[0])
    preds[:,5] = np.round(preds[:,5]).astype(int)
    person = preds[preds[:, 5] == pars['objs'][0]]
    helmet = preds[preds[:, 5] == pars['objs'][1]]
    other = preds[~np.isin(preds[:,5],pars['objs'])]
    # 重塑为(b1数量, 1, 4)和(1, b2数量, 4)以支持广播
    person_coords = person[:, :4].reshape(1, -1, 4)  # 形状: (3, 1, 4)
    helmet_coords = helmet[:, :4].reshape(-1, 1, 4)  # 形状: (1, 2, 4)

    # 批量计算交集坐标
    inter_x1 = np.maximum(person_coords[..., 0], helmet_coords[..., 0])  # 形状: (3, 2)
    inter_y1 = np.maximum(person_coords[..., 1], helmet_coords[..., 1])
    inter_x2 = np.minimum(person_coords[..., 2], helmet_coords[..., 2])
    inter_y2 = np.minimum(person_coords[..., 3], helmet_coords[..., 3])

    # 计算交集面积（确保非负）
    inter_area = np.maximum(0, inter_x2 - inter_x1) * np.maximum(0, inter_y2 - inter_y1)  # 形状: (3, 2)

    # 计算b1和b2各自的面积
    person_area = (person_coords[..., 2] - person_coords[..., 0]) * (person_coords[..., 3] - person_coords[..., 1])  # 形状: (3, 1)
    helmet_area = (helmet_coords[..., 2] - helmet_coords[..., 0]) * (helmet_coords[..., 3] - helmet_coords[..., 1])  # 形状: (1, 2)

    # 计算并集面积和交并比(IoU)
    union_area = person_area + helmet_area - inter_area + 0.001  # 形状: (3, 2)
    iou = inter_area / union_area  # 形状: (3, 2)，每行对应b1元素，每列对应b2元素

    # 找到b2中与任何b1元素IoU>0.25的索引
    # 对每行（b1元素）取最大值，再判断是否有任何b1元素满足条件
    person_mask = np.any(iou > pars['iou'], axis=0)  # 形状: (2,)，True表示符合条件
    # logger.info('异常person_mask:',person_mask)
    # logger.info('异常person:', person)
    inhelmet = person[person_mask]
    unhelmet = person[~person_mask]
    unhelmet[:,5] = pars['unhelmet']
    preds = np.row_stack((inhelmet,helmet,unhelmet,other))


    return [preds.tolist(),'']

if __name__ == "__main__":
    # 对应DJI_20230306140129_0001_Z_165.jpg检测结果
    # preds=[[6.49000e+02, 2.91000e+02, 1.07900e+03, 7.33000e+02, 9.08165e-01, 3.00000e+00],
    #  [8.11000e+02, 2.99000e+02, 1.31200e+03, 7.65000e+02, 8.61268e-01, 3.00000e+00],
    #  [7.05000e+02, 1.96000e+02, 7.19000e+02, 2.62000e+02, 5.66877e-01, 0.00000e+00]]

    # 对应DJI_20230306152702_0001_Z_562.jpg检测结果
    preds = [[7.62000e+02, 7.14000e+02, 1.82800e+03, 9.51000e+02, 9.00902e-01, 3.00000e+00],
             [2.00000e+01, 3.45000e+02, 1.51300e+03, 6.71000e+02, 8.81440e-01, 3.00000e+00],
             [8.35000e+02, 8.16000e+02, 8.53000e+02, 8.30000e+02, 7.07651e-01, 0.00000e+00],
             [1.35600e+03, 4.56000e+02, 1.42800e+03, 4.94000e+02, 6.70549e-01, 2.00000e+00]]
    print('before :\n ', preds)
    # preds=torch.tensor(preds)    #返回的预测结果
    imgwidth = 1920
    imgheight = 1680
    pars = {'imgSize': (imgwidth, imgheight), 'wRation': 1 / 6.0, 'hRation': 1 / 6.0, 'smallId': 0, 'bigId': 3,
            'newId': 4, 'recScale': 1.2}
    # 'smallId':0(国旗)，'bigId':3(船只),wRation和hRation表示判断的阈值条件，newId--新目标的id
    # yyy = channel2_post_process([preds], pars)  # 送入后处理函数
    #
    # print('after :\n ', yyy)