import cv2 as cv
import numpy as np
from re import search
#import os

class yolo:
    def __init__(self, cfg_path, weights_path, classes, confidence_threshold = 0.25):
        #does this matter? maybe not
        #grab width and height from config
        with open(cfg_path) as f:
            cfg = f.read()
            self.height = int(search("height[/s]*=[/s]*([0-9]*)\n", cfg).group(1))
            self.width = int(search("width[/s]*=[/s]*([0-9]*)\n", cfg).group(1))
        self.net = cv.dnn.readNetFromDarknet(cfg_path, weights_path)
        self.net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
        self.classes = classes
        self.confidence_threshold = confidence_threshold

        #palette for drawing on images 
        palette = np.arange(0, 255, dtype=np.uint8).reshape(1, 255, 1)
        palette = cv.applyColorMap(palette, cv.COLORMAP_HSV).squeeze(0)
        np.random.shuffle(palette)
        self.palette = palette
    def show_predictions(self, predictions, input):
        img = input.copy()
        for (x,y,w,h, confidence, label, classID) in predictions:
            #print(label)
           
            cv.rectangle(img, (x, y), (x + w, y + h), tuple(self.palette[classID].tolist()), 2)
            text = "{}: {:.2f}".format(label, confidence)
            cv.putText(img, text, (x, y - 5), cv.FONT_HERSHEY_SIMPLEX, 0.5, tuple(self.palette[classID].tolist()), 1)
            # for (value,key) in zip([name, x,y,w,h, confidence, label, classID], detections.keys()):
            #     detections[key].append(value)
        return img

    def run_inference(self, img, h_offset, v_offset):
        """forward pass on a list of images

            returns: list of predictions in the format
            [(x, y, w, h, conf, class),..] e.g. [ (215, 47, 199, 258, 0.8504674434661865, 'Sporophyte')] 
        """




        ln = self.net.getLayerNames()
        #I don't understand what this is doing
        #fix from https://stackoverflow.com/questions/32978575/how-to-fix-indexerror-invalid-index-to-scalar-variable
        #the original code is from
        #https://opencv-tutorial.readthedocs.io/en/latest/yolo/yolo.html
        ln = [ln[i - 1] for i in self.net.getUnconnectedOutLayers()]

        # construct a blob from the image
        
        
        blob = cv.dnn.blobFromImage(img, 1/255.0, (self.width, self.height), swapRB=True, crop=False)
        self.net.setInput(blob)
        outputs = self.net.forward(ln)
        boxes = []
        confidences = []
        classIDs = []

        #this assumes all images are the same size
        h, w = img.shape[:2]

        for output in outputs:
            for detection in output:
                scores = detection[5:]
                classID = np.argmax(scores)
                confidence = scores[classID]
                #todo: parameterize confidence threshold
                if confidence > self.confidence_threshold:
                    box = detection[:4] * np.array([w, h, w, h])
                    (centerX, centerY, width, height) = box.astype("int")
                    x = int(centerX - (width / 2))
                    y = int(centerY - (height / 2))
                    box = [x, y, int(width), int(height)]
                    boxes.append(box)
                    confidences.append(float(confidence))
                    classIDs.append(classID)
        indices = cv.dnn.NMSBoxes(boxes, confidences, self.confidence_threshold, 0.4)
        outputs = []
        if len(indices) > 0:
            for i in indices.flatten():
                (x, y) = (boxes[i][0], boxes[i][1])
                (w, h) = (boxes[i][2], boxes[i][3])
                label = self.classes[classIDs[i]]
                confidence = confidences[i]
                outputs.append((x + v_offset,y + h_offset,w,h,confidence,label, classIDs[i]))
        return outputs
        """
        colors = np.random.randint(0, 255, size=(len(classes), 3), dtype='uint8')

        if len(indices) > 0:
            for i in indices.flatten():
                (x, y) = (boxes[i][0], boxes[i][1])
                (w, h) = (boxes[i][2], boxes[i][3])
                color = [int(c) for c in colors[classIDs[i]]]
                cv.rectangle(img, (x, y), (x + w, y + h), color, 2)
                text = "{}: {:.4f}".format(classes[classIDs[i]], confidences[i])
                cv.putText(img, text, (x, y - 5), cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
        debug = 1
        cv.imwrite("outputs\\" + imgPath.replace('frames\\',''), img)
        """

#model = yolo("network_parameters\\yolov4-eggs.cfg", "network_parameters\\yolov4-eggs_final.weights",classes)

# files = os.listdir("./frames")
# for frame in files:
#     output = model.run_inference("frames\\" + frame)
"""

vidcap = cv.VideoCapture('C:\\Users\\Adam Honts\\Documents\\Kelp\\5-11\\videos\\220520152231.mp4')
success,image = vidcap.read()
count = 0
while success:
    cv.imwrite("frames\\frame%d.jpg" % count, image)     # save frame as JPEG file
    success,image = vidcap.read()
    print('Read a new frame: ', success)
    count += 1
"""