"""
 Function to extract  Detection and Classification data from the last layer of a propagated cv.dnn_Net network
 input: (ndarray) outputBlobs, (ndarray) frame, (cv.dnn_Net) network,  confThreshold=0.5, nmsThreshold=0.0
 output:
    classIds: 	    list (int)
    confidences:   	list (float)  probability
    boxes: 		    list of  Tuples (x,y,w,h)
"""
import cv2 as cv
import numpy as np


def processDetectionData(outputBlobs, frame, network, confThreshold=0.5, nmsThreshold=0.0):
    frameHeight = frame.shape[0]
    frameWidth = frame.shape[1]

    layerNames = network.getLayerNames()
    lastLayerId = network.getLayerId(layerNames[-1])
    lastLayer = network.getLayer(lastLayerId)

    classIds = []
    confidences = []
    boxes = []
    if lastLayer.type == 'DetectionOutput':
        # Network produces output blob with a shape 1x1xNx7 where N is a number of
        # detections and an every detection is a vector of values
        # [batchId, classId, confidence, left, top, right, bottom]
        for detection in outputBlobs[0, 0]:
            confidence = detection[2]
            left = int(detection[3])
            top = int(detection[4])
            right = int(detection[5])
            bottom = int(detection[6])
            width = right - left + 1
            height = bottom - top + 1
            if width <= 2 or height <= 2:
                left = int(detection[3] * frameWidth)
                top = int(detection[4] * frameHeight)
                right = int(detection[5] * frameWidth)
                bottom = int(detection[6] * frameHeight)
                width = right - left + 1
                height = bottom - top + 1
            if (confidence > confThreshold and
                    left > 0 and top > 0 and
                    left < frameWidth-10 and top < frameHeight-10 and
                    width > 10 and height > 10 ):
                classIds.append(int(detection[1]) - 1)  # Skip background label
                confidences.append(float(confidence))
                boxes.append([left, top, width, height])

    elif lastLayer.type == 'Region':
        # Network produces output blob with a shape NxC where N is a number of
        # detected objects and C is a number of classes + 4 where the first 4
        # numbers are [center_x, center_y, width, height]
        for detection in outputBlobs:
            scores = detection[5:]
            classId = np.argmax(scores)
            confidence = scores[classId]
            width = int(detection[2] * frameWidth)
            height = int(detection[3] * frameHeight)
            left = int(detection[0] * frameWidth - width/2)
            top = int(detection[1] * frameHeight - height/2)
            if (confidence > confThreshold and
                    left > 0 and top > 0 and
                    left < frameWidth-10 and top < frameHeight-10 and
                    width > 10 and height > 10 ):
                classIds.append(classId)
                confidences.append(float(confidence))
                boxes.append([left, top, width, height])
    else:
        print('Unknown output layer type: ' + lastLayer.type)
        return

    # NMS is used inside Region layer only on DNN_BACKEND_OPENCV for another backends we need NMS in sample
    # or NMS is required if number of outputs > 1
    if len(classIds) > 1 and lastLayer.type == 'Region':
        indices = []
        classIds = np.array(classIds)
        boxes = np.array(boxes)
        confidences = np.array(confidences)
        unique_classes = set(classIds)
        for cl in unique_classes:
            class_indices = np.where(classIds == cl)[0]
            conf = confidences[classIds == cl]
            box = boxes[classIds == cl].tolist()
            nms_indices = cv.dnn.NMSBoxes(box, conf, confThreshold, nmsThreshold)
            nms_indices = nms_indices[:] if len(nms_indices)>0 else []
            indices.extend(class_indices[nms_indices])
            classIds = classIds[indices].tolist()
            confidences = confidences[indices].tolist()
            boxes = boxes[indices].tolist()


    return classIds, confidences, boxes