"""

PCB_yolo.py - Main hand detection wrapper for PCB Mask R-CNN network


"""

import cv2
import glob
from main.YOLO.yolo import YOLO

class PCByolo:

    def __init__(self, img_path, model_path):

        self.yolo = YOLO(model_path + "/cross-hands.cfg",
                         model_path + "/cross-hands.weights",
                         ["hand"])

        self.yolo.size = int (416)
        self.yolo.confidence = float(0.2)
        self.files = sorted(glob.glob("%s/*.JPG" % img_path))
        self.time = 0

    def perform_detection(self):

        out = list()
        conf_sum = 0
        time_sum = 0
        detection_count = 0

        for file in self.files:
            mat = cv2.imread(file)

            width, height, inference_time, results = self.yolo.inference(mat)

            print("Processing file " + file.split("/")[-1] + "...")
            print("Found " + str(len(results)) + " objects.")
            print("Processing time: " + str(round(inference_time, 2)) + "s.")

            time_sum += inference_time

            for detection in results:
                id, name, confidence, x, y, w, h = detection
                conf_sum += confidence
                detection_count += 1
                print("Coordinates: X" + str(x) + " Y" + str(y) + " W" + str(w) + " H" + str(h))
                #out.append([file, x, y, w, h])
                out.append([file, y, x, int(y+h), int(x+w)])
                text = "%s (%s)" % (name, round(confidence, 2))

        print("Average confidence: "+ str(round(conf_sum / detection_count, 2 )))

        self.out = out
        self.time = time_sum

    def get_coordinates(self, image):

        res = list()

        for i in self.out:
            print(str(i[0]))
            if(i[0] == image):
                res.append(i[1:])

        return res

    def get_results(self):
        return self.out

    def get_time(self):
        return self.time

    def get_images(self):
        return len(self.files)

