from enum import Enum
import numpy as np

lane_error = 0

class tlabel(Enum):
    TOWER = 0
    SIGN = 1
    SHELTER = 2
    HOSPITAL = 3
    BASKET = 4
    BASE = 5
    YBALL = 6
    SPILLER = 7
    RMARK = 8
    RBLOCK = 9
    RBALL = 10
    MPILLER = 11
    LPILLER = 12
    LMARK = 13
    BBLOCK = 14
    BBALL = 15
'''
description: label_filter 的测试数据
'''
test_resp = {
    'code': 0,
    'data': np.array([
        [4., 0.97192055, 26.64415, 228.26755, 170.16872, 357.6216],
        [4., 0.97049206, 474.0152, 251.2854, 612.91644, 381.6831],
        [5., 0.972649, 250.84174, 238.43622, 378.115, 367.34906]
    ])
}
test1_resp = {
    'code': 0,
    'data': np.array([])
}
'''
description: yolo 目标检测标签过滤器，需要传入连接到 yolo server 的 socket 对象
'''
class label_filter:
    def __init__(self, socket, threshold=0.6):
        self.num = 0
        self.pos = []
        self.socket = socket
        self.threshold = threshold
        
        self.img_size = (320, 240) 
    '''
    description: 向 yolo server 请求目标检测数据
    param {*} self
    return {*}
    '''    
    def get_resp(self):
        self.socket.send_string('')
        response = self.socket.recv_pyobj()
        return response
    '''
    description: 切换 yolo server 视频源 在分叉路口时目标检测需要使用前摄
    param {*} self
    param {*} camera_id 1 或者 2 字符串
    return {*}
    '''    
    def switch_camera(self,camera_id):
        if camera_id == 1 or camera_id == 2:
            self.socket.send_string(f'{camera_id}')
            response = self.socket.recv_pyobj()
            return response
    '''
    description: 对模型推理推理结果使用 threshold 过滤 默认阈值为 0.5
    param {*} self
    param {*} data get_resp 返回的数据
    return {bool,array}
    '''    
    def filter_box(self,data):
        if len(data) > 0:
            expect_boxes = (data[:, 1] > self.threshold) & (data[:, 0] > -1)
            np_boxes = data[expect_boxes, :]
            results = [
                [
                    item[0], # 'label': 
                    item[1], # 'score': 
                    item[2], # 'xmin':
                    item[3], # 'ymin':
                    item[4], # 'xmax':
                    item[5]  # 'ymax':
                ]
                for item in np_boxes
            ]
            if len(results) > 0:
                return True, np.array(results)
        return False, None
    '''
    description: 根据传入的标签过滤，返回该标签的个数以及 box
    param {*} self
    param {*} tlabel
    return {int, array}
    '''    
    def get(self, tlabel):
        # 循环查找匹配的标签值
        # 返回对应标签的个数，以及坐标列表
        response = self.get_resp()
        if response['code'] == 0:
            ret, results = self.filter_box(response['data'])
            if ret:
                expect_boxes = (results[:, 0] == tlabel.value)
                boxes = results[expect_boxes, :]
                self.num = len(boxes)
                self.pos = boxes[:, 2:] # [[x1 y1 x2 y2]]
                return self.num, self.pos
        return 0, []
    '''
    description: 
    param {*} self
    param {*} tlabel_list
    return {*}
    '''    
    def get_mult(self, tlabel_list):
        response = self.get_resp()
        if response['code'] == 0:
            ret, results = self.filter_box(response['data'])
            target_counts = len(tlabel_list)
            counts = 0
            if ret:
                for tlabel in tlabel_list:
                    expect_boxes = (results[:, 0] == tlabel.value)
                    has_true = np.any(expect_boxes)
                    if has_true:
                        counts += 1
                    else:
                        return False, []
                if counts == target_counts:
                    return True, counts
                return False, []
            return False, []    
        return False, []
    '''
    description: 判断传入的标签是否存在，存在返回 True
    param {*} self
    param {*} tlabel
    return {bool}
    '''    
    def find(self, tlabel):
        response = self.get_resp()
        if response['code'] == 0:
            ret, results = self.filter_box(response['data'])
            if ret:
                expect_boxes = (results[:, 0] == tlabel.value)
                boxes = results[expect_boxes, :]
                if len(boxes) != 0:
                    return True 
        return False
    '''
    description: 根据传入的标签，
    param {*} self
    param {*} tlabel
    return {*}
    '''    
    def aim_left(self, tlabel):
        # 如果标签存在，则返回列表中位置最靠左的目标框和中心的偏移值
        response = self.get_resp()
        if response['code'] == 0:
            ret, results = self.filter_box(response['data'])
            if ret:
                expect_boxes = (results[:, 0] == tlabel.value)
                boxes = results[expect_boxes, :]
                if len(boxes) == 0:
                    return (False, )
                xmin_values = boxes[:, 2] # xmin
                xmin_index = np.argmin(xmin_values)
                error = (boxes[xmin_index][4] + boxes[xmin_index][2] - self.img_size[0]) / 2
                return (True, error)
        return (False, )
    def aim_right(self, tlabel):
        # 如果标签存在，则返回列表中位置最靠右的目标框和中心的偏移值
        response = self.get_resp()
        if response['code'] == 0:
            ret, results = self.filter_box(response['data'])
            if ret:
                expect_boxes = (results[:, 0] == tlabel.value)
                boxes = results[expect_boxes, :]
                if len(boxes) == 0:
                    return (False, )
                xmax_values = boxes[:, 4] # xmax
                xmax_index = np.argmax(xmax_values)
                error = (boxes[xmax_index][4] + boxes[xmax_index][2] - self.img_size[0]) / 2
                return (True, error)
        return (False, )
    def aim_near(self, tlabel):
        # 如果标签存在，则返回列表中位置最近的目标框和中心的偏移值
        response = self.get_resp()
        if response['code'] == 0:
            ret, results = self.filter_box(response['data'])
            if ret:
                expect_boxes = (results[:, 0] == tlabel.value)
                boxes = results[expect_boxes, :]
                if len(boxes) == 0:
                    return (False, 0)
                center_x_values = np.abs(boxes[:, 2] + boxes[:, 4] - self.img_size[0])
                center_x_index = np.argmin(center_x_values)
                error = (boxes[center_x_index][4] + boxes[center_x_index][2] - self.img_size[0]) / 2
                return (True, error)
        return (False, 0)

# class Calibrate:
#     def __init__(self,by_cmd):
#         # 车控制对象初始化
#         self.by_cmd = by_cmd
#     def aim(self,error):
#         self.by_cmd.send_distance_x(error,)


if __name__ == '__main__':
    obj = label_filter(None)
    # results = obj.filter_box(resp['data'])
    # expect_boxes = (results[:, 0] == tlabel.SPILLAR.value)
    # np_boxes = results[expect_boxes, :]
    # print(np_boxes[:, 2:])
    # print(len(np_boxes))
    print(obj.find(tlabel.BBALL))
    print(obj.aim_left(tlabel.BBALL))
    print(obj.aim_right(tlabel.BBALL))
    print(obj.aim_near(tlabel.BBALL))
    print(obj.get(tlabel.HOSPITAL))