import sys
import cv2
import numpy

sys.dont_write_bytecode = True

# Source
from helper import onnxSessionBuild

pathModel = "./RT-DETR-L_wireless_table_cell_det/"

imageSizeModel = 640

scoreThreshold = 0.3
levelBoxNms = 0.5
levelBoxContained = 0.9
countContainedMinimum = 2

onnxSession = onnxSessionBuild(f"{pathModel}onnx/rt-detr-l_wireless_table_cell_det.onnx")

def overlapCalculate(coordinateFirstList, coordinateSecondList):
    areaFirst = (coordinateFirstList[2] - coordinateFirstList[0]) * (coordinateFirstList[3] - coordinateFirstList[1])
    areaSecond = (coordinateSecondList[2] - coordinateSecondList[0]) * (coordinateSecondList[3] - coordinateSecondList[1])

    x1 = max(coordinateFirstList[0], coordinateSecondList[0])
    y1 = max(coordinateFirstList[1], coordinateSecondList[1])
    x2 = min(coordinateFirstList[2], coordinateSecondList[2])
    y2 = min(coordinateFirstList[3], coordinateSecondList[3])

    if x2 <= x1 or y2 <= y1:
        return {"intersectionOverUnion": 0.0, "containment": 0.0, "areaFirst": areaFirst, "areaSecond": areaSecond}

    areaIntersection = (x2 - x1) * (y2 - y1)

    return {
        "intersectionOverUnion": areaIntersection / float(areaFirst + areaSecond - areaIntersection),
        "containment": areaIntersection / float(min(areaFirst, areaSecond)),
        "areaFirst": areaFirst,
        "areaSecond": areaSecond
    }

def suppressionNonMaximum(itemList):
    resultList = []

    itemSortedList = sorted(itemList, key=lambda itemObject: itemObject["score"], reverse=True)

    for a in range(len(itemSortedList)):
        isKeep = True

        for b in range(len(resultList)):
            overlapObject = overlapCalculate(itemSortedList[a]["coordinate"], resultList[b]["coordinate"])

            if overlapObject["intersectionOverUnion"] >= levelBoxNms:
                isKeep = False

                break

        if isKeep:
            resultList.append(itemSortedList[a])

    return resultList

def suppressionContainer(itemList):
    resultList = []

    for a in range(len(itemList)):
        countContained = 0

        for b in range(len(itemList)):
            if a == b:
                continue

            overlapObject = overlapCalculate(itemList[a]["coordinate"], itemList[b]["coordinate"])

            if overlapObject["containment"] >= levelBoxContained and overlapObject["areaFirst"] > overlapObject["areaSecond"]:
                countContained += 1

        if countContained < countContainedMinimum:
            resultList.append(itemList[a])

    return resultList

def inference(imageRgb):
    resultList = []

    imageHeight, imageWidth = imageRgb.shape[0:2]
    imageResized = cv2.resize(imageRgb, (imageSizeModel, imageSizeModel), interpolation=cv2.INTER_CUBIC).astype(numpy.float32) / 255.0

    tensor = numpy.expand_dims(imageResized.transpose((2, 0, 1)), axis=0).astype(numpy.float32)

    tensorFeedObject = {
        "image": tensor,
        "im_shape": numpy.array([[imageSizeModel, imageSizeModel]], dtype=numpy.float32),
        "scale_factor": numpy.array([[imageSizeModel / float(imageHeight), imageSizeModel / float(imageWidth)]], dtype=numpy.float32)
    }

    tensorOutputList = onnxSession.run(None, tensorFeedObject)

    boxCount = int(tensorOutputList[1][0]) if len(tensorOutputList) > 1 else len(tensorOutputList[0])

    for a in range(boxCount):
        value = tensorOutputList[0][a]

        score = float(value[1])
        x1 = max(0.0, min(float(value[2]), float(imageWidth)))
        y1 = max(0.0, min(float(value[3]), float(imageHeight)))
        x2 = max(0.0, min(float(value[4]), float(imageWidth)))
        y2 = max(0.0, min(float(value[5]), float(imageHeight)))

        if score >= scoreThreshold and x2 > x1 and y2 > y1:
            resultList.append({
                "score": score,
                "coordinate": [x1, y1, x2, y2]
            })

    resultList = suppressionNonMaximum(resultList)
    resultList = suppressionContainer(resultList)

    return resultList

image = cv2.imread(sys.argv[1])
imageRgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

itemList = inference(imageRgb)

for a in range(len(itemList)):
    coordinateList = []

    for b in range(len(itemList[a]["coordinate"])):
        coordinateList.append(int(round(itemList[a]["coordinate"][b])))

    print(f"{itemList[a]['score']:.6f} | {coordinateList}")