import cv2
import os
import fnmatch
import time
import torch
from random import randrange

# custom py implementation of Resnet50 model
import ResNet_imp
from ResNet_imp import init_model
from ResNet_imp import get_prediction

# mode; 0 - calculate, 1 - show visually
mode = 0
# skip n frames between prediction
skip = 8
# crop ; 0 - resize img to 260x260 and perform random 224x224 crop ;  crop: 1 - precise crop
precise_crop = 0


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


# draw rectangular around object
def draw_box(f_img, f_resolution1, f_resolution2, f_x, f_y, f_w, f_h, f_class):

    # calculating bb to usable format
    usable_x = int(round(f_resolution1*float(f_x)))
    usable_y = int(round(f_resolution2*float(f_y)))
    usable_w_half = int(round(f_resolution1*float(f_w))/2)
    usable_h_half = int(round(f_resolution2*float(f_h))/2)

    cv2.rectangle(f_img, (max(0, usable_x-usable_w_half), max(0, usable_y-usable_h_half)),
                  (min(f_resolution1, usable_x+usable_w_half), min(f_resolution2, usable_y+usable_h_half)),
                  (255, 0, 255), 3, 1)
    if float(f_y) > 0.5:
        cv2.putText(f_img, f_class, (usable_x, usable_y-usable_h_half-10), cv2.FONT_HERSHEY_DUPLEX, 0.7, (0, 255, 0), 2)
    else:
        cv2.putText(f_img, f_class, (usable_x, usable_y+usable_h_half+25), cv2.FONT_HERSHEY_DUPLEX, 0.7, (0, 255, 0), 2)


def crop_exact(f_img, f_resolution1, f_resolution2, f_x, f_y, f_w, f_h,):
    # calculating bb to usable format
    usable_x = int(round(f_resolution1 * float(f_x)))
    usable_y = int(round(f_resolution2 * float(f_y)))
    usable_w_half = int(round(f_resolution1 * float(f_w)) / 2)
    usable_h_half = int(round(f_resolution2 * float(f_h)) / 2)

    return f_img[max(0, usable_y - usable_h_half):min(f_resolution2, usable_y + usable_h_half),
                 max(0, usable_x - usable_w_half):min(f_resolution1, usable_x + usable_w_half)]


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

# device selection
device = 'cpu'
if torch.cuda.is_available():
    device = 'cuda'
# init ResNet_model ###########
model_main = init_model(device)

number_of_frames = [0, 0, 0, 0, 0, 0]
number_of_true_positive = [0, 0, 0, 0, 0, 0]
number_of_false_positive = [0, 0, 0, 0, 0, 0]
number_of_true_negative = [0, 0, 0, 0, 0, 0]
number_of_false_negative = [0, 0, 0, 0, 0, 0]

end_testing = 0
for i in range(len(videos_array)):
    # 'ESC'
    if end_testing == 1:
        break

    # path join
    video_path_plus_name = str(videos_path) + '\\' + str(videos_array[i])
    label_path_plus_name = str(labels_path) + '\\' + str(labels_array[i])

    # source of video
    cap = cv2.VideoCapture(video_path_plus_name)

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

    while True:
        # skip n frames from video
        video_end = 0
        for j in range(skip):
            ret, _ = cap.read()
            if ret == 0:
                video_end = 1
                break
            # 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()

        if video_end == 1:
            break

        # get frame
        ret, frame = cap.read()
        if ret == 0:
            break

        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        # read label part for current frame
        string = f.readline()
        img_cropped_224 = []
        prediction = []
        while string != "cap\n" and string != "":
            # line split
            bb = string.split(" ")

            # add cropped image to array (exact on bb edge - no padding)
            img_cropped = crop_exact(frame, width, height, bb[1], bb[2], bb[3], bb[4])
            # resize
            if precise_crop == 0:
                img_cropped_260 = cv2.resize(img_cropped, (260, 260))
                x_shift = randrange(0, 36)
                y_shift = randrange(0, 36)
                img_cropped_224.append(img_cropped_260[0+y_shift:224+y_shift, 0+x_shift:224+x_shift])
            else:
                img_cropped_224.append(cv2.resize(img_cropped, (224, 224)))
            # prediction of bb
            pred = get_prediction([img_cropped_224[-1]], model_main, device)
            prediction.append(pred[0])

            # draw bb to origin frame
            if mode == 1:
                draw_box(frame, width, height, bb[1], bb[2], bb[3], bb[4], bb[0])
            # read new line

            # number of predictions
            number_of_frames[int(bb[0])] = number_of_frames[int(bb[0])] + 1
            # number of successfully predicted

            # TN
            for j in range(len(number_of_true_negative)):
                if (j != int(bb[0])) and (j != int(prediction[-1])):
                    number_of_true_negative[j] += 1

            if int(prediction[-1]) == int(bb[0]):
                # TP
                number_of_true_positive[int(bb[0])] += 1
            else:
                # FP
                number_of_false_positive[int(prediction[-1])] += 1
                # FN
                number_of_false_negative[int(bb[0])] += 1

            # read new line
            string = f.readline()

        if mode == 1:
            # display capture
            cv2.imshow("Video stream", frame)
            # display crop (only first bb in frame)
            if len(img_cropped_224) > 0:
                # put predicted class to cropped image (only first bb in frame)
                cv2.putText(img_cropped_224[0], str(prediction[0]), (20, 20), cv2.FONT_HERSHEY_DUPLEX, 0.7,
                            (0, 255, 0), 2)
                # display crop (only first bb in frame)
                cv2.imshow("Cropped by label", img_cropped_224[0])

            time.sleep(0.2)

        if mode == 0:
            continue
        key = cv2.waitKey(30)
        # 'q' next video
        if key == 113:
            break
        # 'ESC' stop testing
        if key == 27:
            end_testing = 1
            break

    f.close()

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

F1 = []
for i in range(len(labels)):
    print("class: "+str(labels[i]))
    print("number of images "+str(number_of_frames[i]))
    print("TP: "+str(number_of_true_positive[i]))
    print("FP: "+str(number_of_false_positive[i]))
    print("FN: "+str(number_of_false_negative[i]))
    print("TN: "+str(number_of_true_negative[i]))
    # calculating F1
    if number_of_true_positive[i] > 0:
        F1.append(number_of_true_positive[i] / (
                  number_of_true_positive[i] + 0.5 * (number_of_false_positive[i] + number_of_false_negative[i])))
    else:
        F1.append(0)

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


