import cv2
import os
import fnmatch
import time
import torch

# modular program iota
import model_iota_function
from model_iota_function import model_iota

# 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 testing videos
videos_path = r""
videos_array = listdir(videos_path)
# get labels of testing videos
labels_path = \
    r""
labels_array = listdir(labels_path)

valid = 1
if len(videos_array) == len(labels_array):
    for i in range(len(videos_array)):
        name1 = videos_array[i][:-4]
        name2 = labels_array[i][:-4]
        if name1 != name2:
            valid = 0
            break

if valid == 0:
    exit()


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(videos_array)):
    print(i)
    # path join
    video_path_plus_name = str(videos_path) + '\\' + str(videos_array[i])
    label_path_plus_name = str(labels_path) + '\\' + str(labels_array[i])

    # model for prediction for whole video
    prediction_output = model_iota(video_path_plus_name)

    # open label and check label
    f = open(label_path_plus_name, 'r')
    string = f.readline()
    if string != "cap\n":
        continue

    if mode == 1:
        # camera
        cv2.namedWindow("Frame")
        cv2.moveWindow("Frame", 40, 25)
        cap = cv2.VideoCapture(video_path_plus_name)

    print(len(prediction_output))
    # print(prediction_output)
    for j in range(len(prediction_output)):
        if mode == 1:
            ret, frame = cap.read()
            if ret == 0:
                break

        # skip if output is [] (without detection)
        if len(prediction_output[j]) == 0:
            # skip labels (otherwise labels would be de-synchronize)
            string = f.readline()
            while string != "cap\n" and 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()

        # skip if output is [-1] (not prediction for frame)
        elif prediction_output[j][0] == -1:
            # skip labels (otherwise labels would be de-synchronize)
            string = f.readline()
            while string != "cap\n" and string != "":
                # line split
                bb = string.split(" ")
                # read new line
                string = f.readline()
        else:
            # read prediction for this frame
            frame_prediction = prediction_output[j]
            pred_array = []
            for k in range(len(frame_prediction)):
                pred_array.append(frame_prediction[k])

            # print("pred: "+str(pred_array))

            # 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()

            # print("yolo: " + str(bb_array))

            # accuracy of prediction (comparing with yolo bb)
            num_of_false_detection, colliding_array = colliding_bb(bb_array, pred_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
                    number_of_multiple_detection_of_object[int(bb[0])] += (len(index_of_pred) - 1)
                    for n in range(len(index_of_pred)):
                        pred = pred_array[int(index_of_pred[n])]

                        # 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)

            # visualization
                # draw bb boxes from labels and prediction
            if mode == 1:
                for k in range(len(pred_array)):
                    # prediction extraction from tuple
                    pred_value, pred_x1, pred_y1, pred_x2, pred_y2 = pred_array[k]

                    # rectangular of prediction
                    cv2.rectangle(frame, (pred_x1, pred_y1), (pred_x2, pred_y2), (0, 0, 255), 3)
                    # class prediction
                    if float(pred_y1) > 256:
                        cv2.putText(frame, str(pred_value), (pred_x1, pred_y1 - 10), cv2.FONT_HERSHEY_DUPLEX, 0.7,
                                    (0, 255, 0), 2)
                    else:
                        cv2.putText(frame, str(pred_value), (pred_x1, pred_y2 + 25), cv2.FONT_HERSHEY_DUPLEX, 0.7,
                                    (0, 255, 0), 2)

                for k in range(len(bb_array)):
                    # prediction extraction from tuple
                    bb_value, bb_x1, bb_y1, bb_x2, bb_y2 = bb_array[k]

                    # rectangular of prediction
                    cv2.rectangle(frame, (bb_x1, bb_y1), (bb_x2, bb_y2), (0, 255, 0), 3)
                    # class prediction
                    if float(bb_y1) > 256:
                        cv2.putText(frame, str(bb_value), (bb_x1, bb_y1 - 10), cv2.FONT_HERSHEY_DUPLEX, 0.7,
                                    (0, 255, 0), 2)
                    else:
                        cv2.putText(frame, str(bb_value), (bb_x1, bb_y2 + 25), cv2.FONT_HERSHEY_DUPLEX, 0.7,
                                    (0, 255, 0), 2)

                cv2.imshow("Frame", frame)

                cv2.waitKey(1000)

    print("number_of_object")
    print(number_of_object)

    if mode == 1:
        cv2.destroyWindow("Frame")

    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("Successfully detected: "+str(number_of_successfully_detected[i]))
    print("Redundant detection: "+str(number_of_multiple_detection_of_object[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("")
