import cv2
import os
import fnmatch
import time
import torch


# SHOW FRAMES ##########################################################################################################
# mode = 0 - just compute ; mode = 1 - compute + show frames
mode = 0


# get all files of defined type from folder
def listdir(dir_name, pattern="*"):
    return fnmatch.filter(os.listdir(dir_name), pattern)


def colliding_bb(bb_a, pred_a):
    colliding = []
    false_detection_counter = [0, 0, 0, 0, 0, 0]
    for c in range(len(bb_a)):
        _, f_bb_x1, f_bb_y1, f_bb_x2, f_bb_y2 = bb_a[c]
        c_colliding = []
        for v in range(len(pred_a)):
            f_pred_class, f_pred_x1, f_pred_y1, f_pred_x2, f_pred_y2 = pred_a[v]
            # for colliding
            if ((f_pred_x1 <= f_bb_x1) and (f_pred_x2 >= f_bb_x1)) or (
                    (f_pred_x1 <= f_bb_x2) and (f_pred_x2 >= f_bb_x2)) or (
                    (f_pred_x1 > f_bb_x1) and (f_pred_x2 < f_bb_x2)) or (    # box1 in box2
                    (f_pred_x1 < f_bb_x1) and (f_pred_x2 > f_bb_x2)):        # box2 in box1
                if ((f_pred_y1 <= f_bb_y1) and (f_pred_y2 >= f_bb_y1)) or (
                        (f_pred_y1 <= f_bb_y2) and (f_pred_y2 >= f_bb_y2)) or (
                        (f_pred_y1 > f_bb_y1) and (f_pred_y2 < f_bb_y2)) or (  # box1 in box2
                        (f_pred_y1 < f_bb_y1) and (f_pred_y2 > f_bb_y2)):      # box2 in box1
                    c_colliding.append(v)
        if (len(c_colliding) == 0) and (int(f_pred_class) != 5):
            false_detection_counter[int(f_pred_class)] += 1
        colliding.append(c_colliding)

    return false_detection_counter, colliding


# get labels of testing images
labels_path = r""
labels_array = listdir(labels_path)

# prediction of yolo
YOLO_detect_path = r""
YOLO_array = listdir(YOLO_detect_path)

number_of_object = [0, 0, 0, 0, 0, 0]
prediction_TP = [0, 0, 0, 0, 0, 0]
prediction_FP = [0, 0, 0, 0, 0, 0]
prediction_FN = [0, 0, 0, 0, 0, 0]
prediction_TN = [0, 0, 0, 0, 0, 0]

number_of_successfully_detected = [0, 0, 0, 0, 0, 0]
number_of_false_detection = [0, 0, 0, 0, 0, 0]
number_of_object_without_detection = [0, 0, 0, 0, 0, 0]
number_of_multiple_detection_of_object = [0, 0, 0, 0, 0, 0]

for i in range(len(labels_array)):
    print(i)
    # path join
    label_path_plus_name = str(labels_path) + '\\' + str(labels_array[i])
    if labels_array[i] in YOLO_array:
        YOLO_path_plus_name = str(YOLO_detect_path) + '\\' + str(labels_array[i])
    else:
        YOLO_path_plus_name = None

    f = open(label_path_plus_name, 'r')

    # skip if YOLO not predicting
    if YOLO_path_plus_name is None:
        string = f.readline()
        while string != "":
            # line split
            bb = string.split(" ")
            number_of_object_without_detection[int(bb[0])] = number_of_object_without_detection[int(bb[0])] + 1
            number_of_object[int(bb[0])] = number_of_object[int(bb[0])] + 1
            # read new line
            string = f.readline()

    else:
        f_yolo = open(YOLO_path_plus_name, 'r')
        # read prediction for this frame

        # yolo bb converting to same format as prediction
        bb_array = []
        string = f.readline()
        while string != "cap\n" and string != "":
            # line split
            yolo_bb = string.split(" ")
            usable_x = int(round(640 * float(yolo_bb[1])))
            usable_y = int(round(512 * float(yolo_bb[2])))
            usable_w_half = int(round(640 * float(yolo_bb[3])) / 2)
            usable_h_half = int(round(512 * float(yolo_bb[4])) / 2)
            x1_bb = max(0, usable_x - usable_w_half)
            x2_bb = min(640, usable_x + usable_w_half)
            y1_bb = max(0, usable_y - usable_h_half)
            y2_bb = min(512, usable_y + usable_h_half)
            bb_array.append((yolo_bb[0], x1_bb, y1_bb, x2_bb, y2_bb))

            # read new line
            string = f.readline()

        bb_yolo_array = []
        string = f_yolo.readline()
        while string != "cap\n" and string != "":
            # line split
            yolo_bb = string.split(" ")
            usable_x = int(round(640 * float(yolo_bb[1])))
            usable_y = int(round(512 * float(yolo_bb[2])))
            usable_w_half = int(round(640 * float(yolo_bb[3])) / 2)
            usable_h_half = int(round(512 * float(yolo_bb[4])) / 2)
            x1_bb = max(0, usable_x - usable_w_half)
            x2_bb = min(640, usable_x + usable_w_half)
            y1_bb = max(0, usable_y - usable_h_half)
            y2_bb = min(512, usable_y + usable_h_half)
            bb_yolo_array.append((yolo_bb[0], x1_bb, y1_bb, x2_bb, y2_bb))

            # read new line
            string = f_yolo.readline()

        f_yolo.close()
        # print("yolo: " + str(bb_array))

        # accuracy of prediction (comparing with yolo bb)
        num_of_false_detection, colliding_array = colliding_bb(bb_array, bb_yolo_array)
        # number of false detection
        for m in range(len(num_of_false_detection)):
            number_of_false_detection[m] += num_of_false_detection[m]
        for m in range(len(colliding_array)):
            # number of objects
            bb = bb_array[m]
            number_of_object[int(bb[0])] = number_of_object[int(bb[0])] + 1

            if len(colliding_array[m]) == 0:
                number_of_object_without_detection[int(bb[0])] = number_of_object_without_detection[int(bb[0])] + 1
            else:
                index_of_pred = colliding_array[m]
                number_of_successfully_detected[int(bb[0])] += 1

                for n in range(len(index_of_pred)):
                    pred = bb_yolo_array[int(index_of_pred[n])]
                    number_of_multiple_detection_of_object[int(bb[0])] += (len(index_of_pred) - 1)
                    # TN
                    for r in range(len(prediction_TN)):
                        if r != int(bb[0]) and r != int(pred[0]):
                            prediction_TN[r] += 1

                    if int(pred[0]) == int(bb[0]):
                        # TP
                        prediction_TP[int(bb[0])] += 1
                    else:
                        # FP
                        prediction_FP[int(pred[0])] += 1
                        # FN
                        prediction_FN[int(bb[0])] += 1
        # print(colliding_array)

    print("number_of_object")
    print(number_of_object)

    f.close()

labels = ["Airplane", "Bird", "Drone", "Helicopter", "Insect", "Background"]

# detection result
F1_det = []
for i in range(len(labels)):
    print("class: "+str(labels[i]))
    print("number of objects for detection "+str(number_of_object[i]))
    print("Redundant detection: " + str(number_of_multiple_detection_of_object[i]))
    print("Successfully detected: "+str(number_of_successfully_detected[i]))
    print("False detection: "+str(number_of_false_detection[i]))
    print("Missed: "+str(number_of_object_without_detection[i]))

    # calculating F1
    if number_of_successfully_detected[i] > 0:
        F1_det.append(number_of_successfully_detected[i] / (
                  number_of_successfully_detected[i] +
                  0.5 * (number_of_object_without_detection[i] + number_of_false_detection[i])))
    else:
        F1_det.append(0)

    print("F1 score: "+str(round(F1_det[i]*100, 2))+"%")
    print("")


# prediction results
F1 = []
for i in range(len(labels)):
    print("class: "+str(labels[i]))
    print("number of images "+str(number_of_successfully_detected[i]))
    print("TP: "+str(prediction_TP[i]))
    print("FP: "+str(prediction_FP[i]))
    print("FN: "+str(prediction_FN[i]))
    print("TN: "+str(prediction_TN[i]))
    # calculating F1
    if prediction_TP[i] > 0:
        F1.append(prediction_TP[i] / (
                  prediction_TP[i] + 0.5 * (prediction_FP[i] + prediction_FN[i])))
    else:
        F1.append(0)

    print("F1 score: "+str(round(F1[i]*100, 2))+"%")
    print("")
