#!python3
"""
Python 3 wrapper for identifying objects in images

Requires DLL compilation

Both the GPU and no-GPU version should be compiled; the no-GPU version should be renamed "yolo_cpp_dll_nogpu.dll".

On a GPU system, you can force CPU evaluation by any of:

- Set global variable DARKNET_FORCE_CPU to True
- Set environment variable CUDA_VISIBLE_DEVICES to -1
- Set environment variable "FORCE_CPU" to "true"


To use, either run performDetect() after import, or modify the end of this file.

See the docstring of performDetect() for parameters.

Directly viewing or returning bounding-boxed images requires scikit-image to be installed (`pip install scikit-image`)


Original *nix 2.7: https://github.com/pjreddie/darknet/blob/0f110834f4e18b30d5f101bf8f1724c34b7b83db/python/darknet.py
Windows Python 2.7 version: https://github.com/AlexeyAB/darknet/blob/fc496d52bf22a0bb257300d3c79be9cd80e722cb/build/darknet/x64/darknet.py

@author: Philip Kahn
@date: 20180503
"""
#pylint: disable=R, W0401, W0614, W0703
from ctypes import *
import math
import random
import os

def sample(probs):
    s = sum(probs)
    probs = [a/s for a in probs]
    r = random.uniform(0, 1)
    for i in range(len(probs)):
        r = r - probs[i]
        if r <= 0:
            return i
    return len(probs)-1

def c_array(ctype, values):
    arr = (ctype*len(values))()
    arr[:] = values
    return arr

class BOX(Structure):
    _fields_ = [("x", c_float),
                ("y", c_float),
                ("w", c_float),
                ("h", c_float)]

class DETECTION(Structure):
    _fields_ = [("bbox", BOX),
                ("classes", c_int),
                ("prob", POINTER(c_float)),
                ("mask", POINTER(c_float)),
                ("objectness", c_float),
                ("sort_class", c_int),
                ("uc", POINTER(c_float)),
                ("points", c_int)]

class DETNUMPAIR(Structure):
    _fields_ = [("num", c_int),
                ("dets", POINTER(DETECTION))]

class IMAGE(Structure):
    _fields_ = [("w", c_int),
                ("h", c_int),
                ("c", c_int),
                ("data", POINTER(c_float))]

class METADATA(Structure):
    _fields_ = [("classes", c_int),
                ("names", POINTER(c_char_p))]



#lib = CDLL("/home/pjreddie/documents/darknet/libdarknet.so", RTLD_GLOBAL)
#lib = CDLL("libdarknet.so", RTLD_GLOBAL)
hasGPU = True
cwd = os.path.dirname(__file__)
if os.name == "nt":    
    os.environ['PATH'] = cwd + ';' + os.environ['PATH']
    winGPUdll = os.path.join(cwd, "yolo_cpp_dll.dll")
    winNoGPUdll = os.path.join(cwd, "yolo_cpp_dll_nogpu.dll")
    envKeys = list()
    for k, v in os.environ.items():
        envKeys.append(k)
    try:
        try:
            tmp = os.environ["FORCE_CPU"].lower()
            if tmp in ["1", "true", "yes", "on"]:
                raise ValueError("ForceCPU")
            else:
                print("Flag value '"+tmp+"' not forcing CPU mode")
        except KeyError:
            # We never set the flag
            if 'CUDA_VISIBLE_DEVICES' in envKeys:
                if int(os.environ['CUDA_VISIBLE_DEVICES']) < 0:
                    raise ValueError("ForceCPU")
            try:
                global DARKNET_FORCE_CPU
                if DARKNET_FORCE_CPU:
                    raise ValueError("ForceCPU")
            except NameError:
                pass
            # print(os.environ.keys())
            # print("FORCE_CPU flag undefined, proceeding with GPU")
        if not os.path.exists(winGPUdll):
            raise ValueError("NoDLL")
        lib = CDLL(winGPUdll, RTLD_GLOBAL)
    except (KeyError, ValueError):
        hasGPU = False
        if os.path.exists(winNoGPUdll):
            lib = CDLL(winNoGPUdll, RTLD_GLOBAL)
            print("Notice: CPU-only mode")
        else:
            # Try the other way, in case no_gpu was
            # compile but not renamed
            lib = CDLL(winGPUdll, RTLD_GLOBAL)
            print("Environment variables indicated a CPU run, but we didn't find `"+winNoGPUdll+"`. Trying a GPU run anyway.")
else:
    try:
        lib = CDLL(os.path.join(cwd, "libdarknet.so"), RTLD_GLOBAL)
    except:
        lib=None
if(lib):
    lib.network_width.argtypes = [c_void_p]
    lib.network_width.restype = c_int
    lib.network_height.argtypes = [c_void_p]
    lib.network_height.restype = c_int

    copy_image_from_bytes = lib.copy_image_from_bytes
    copy_image_from_bytes.argtypes = [IMAGE,c_char_p]

    def network_width(net):
        return lib.network_width(net)

    def network_height(net):
        return lib.network_height(net)

    predict = lib.network_predict_ptr
    predict.argtypes = [c_void_p, POINTER(c_float)]
    predict.restype = POINTER(c_float)

    if hasGPU:
        set_gpu = lib.cuda_set_device
        set_gpu.argtypes = [c_int]
        
    init_cpu = lib.init_cpu

    make_image = lib.make_image
    make_image.argtypes = [c_int, c_int, c_int]
    make_image.restype = IMAGE

    get_network_boxes = lib.get_network_boxes
    get_network_boxes.argtypes = [c_void_p, c_int, c_int, c_float, c_float, POINTER(c_int), c_int, POINTER(c_int), c_int]
    get_network_boxes.restype = POINTER(DETECTION)

    make_network_boxes = lib.make_network_boxes
    make_network_boxes.argtypes = [c_void_p]
    make_network_boxes.restype = POINTER(DETECTION)

    free_detections = lib.free_detections
    free_detections.argtypes = [POINTER(DETECTION), c_int]

    free_batch_detections = lib.free_batch_detections
    free_batch_detections.argtypes = [POINTER(DETNUMPAIR), c_int]

    free_ptrs = lib.free_ptrs
    free_ptrs.argtypes = [POINTER(c_void_p), c_int]

    network_predict = lib.network_predict_ptr
    network_predict.argtypes = [c_void_p, POINTER(c_float)]

    reset_rnn = lib.reset_rnn
    reset_rnn.argtypes = [c_void_p]

    load_net = lib.load_network
    load_net.argtypes = [c_char_p, c_char_p, c_int]
    load_net.restype = c_void_p

    load_net_custom = lib.load_network_custom
    load_net_custom.argtypes = [c_char_p, c_char_p, c_int, c_int]
    load_net_custom.restype = c_void_p

    do_nms_obj = lib.do_nms_obj
    do_nms_obj.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]

    do_nms_sort = lib.do_nms_sort
    do_nms_sort.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]

    free_image = lib.free_image
    free_image.argtypes = [IMAGE]

    letterbox_image = lib.letterbox_image
    letterbox_image.argtypes = [IMAGE, c_int, c_int]
    letterbox_image.restype = IMAGE

    load_meta = lib.get_metadata
    lib.get_metadata.argtypes = [c_char_p]
    lib.get_metadata.restype = METADATA

    load_image = lib.load_image_color
    load_image.argtypes = [c_char_p, c_int, c_int]
    load_image.restype = IMAGE

    rgbgr_image = lib.rgbgr_image
    rgbgr_image.argtypes = [IMAGE]

    predict_image = lib.network_predict_image
    predict_image.argtypes = [c_void_p, IMAGE]
    predict_image.restype = POINTER(c_float)

    predict_image_letterbox = lib.network_predict_image_letterbox
    predict_image_letterbox.argtypes = [c_void_p, IMAGE]
    predict_image_letterbox.restype = POINTER(c_float)

    network_predict_batch = lib.network_predict_batch
    network_predict_batch.argtypes = [c_void_p, IMAGE, c_int, c_int, c_int,
                                       c_float, c_float, POINTER(c_int), c_int, c_int]
    network_predict_batch.restype = POINTER(DETNUMPAIR)

    def array_to_image(arr):
        import numpy as np
        # need to return old values to avoid python freeing memory
        arr = arr.transpose(2,0,1)
        c = arr.shape[0]
        h = arr.shape[1]
        w = arr.shape[2]
        arr = np.ascontiguousarray(arr.flat, dtype=np.float32) / 255.0
        data = arr.ctypes.data_as(POINTER(c_float))
        im = IMAGE(w,h,c,data)
        return im, arr

    def classify(net, meta, im):
        out = predict_image(net, im)
        res = []
        for i in range(meta.classes):
            if altNames is None:
                nameTag = meta.names[i]
            else:
                nameTag = altNames[i]
            res.append((nameTag, out[i]))
        res = sorted(res, key=lambda x: -x[1])
        return res

    def detect(net, meta, image, thresh=.5, hier_thresh=.5, nms=.45, debug= False):
        """
        Performs the meat of the detection
        """
        #pylint: disable= C0321
        im = load_image(image, 0, 0)
        if debug: print("Loaded image")
        ret = detect_image(net, meta, im, thresh, hier_thresh, nms, debug)
        free_image(im)
        if debug: print("freed image")
        return ret

    def detect_image(net, meta, im, thresh=.5, hier_thresh=.5, nms=.45, debug= False):
        #import cv2
        #custom_image_bgr = cv2.imread(image) # use: detect(,,imagePath,)
        #custom_image = cv2.cvtColor(custom_image_bgr, cv2.COLOR_BGR2RGB)
        #custom_image = cv2.resize(custom_image,(lib.network_width(net), lib.network_height(net)), interpolation = cv2.INTER_LINEAR)
        #import scipy.misc
        #custom_image = scipy.misc.imread(image)
        #im, arr = array_to_image(custom_image)		# you should comment line below: free_image(im)
        num = c_int(0)
        if debug: print("Assigned num")
        pnum = pointer(num)
        if debug: print("Assigned pnum")
        predict_image(net, im)
        letter_box = 0
        #predict_image_letterbox(net, im)
        #letter_box = 1
        if debug: print("did prediction")
        #dets = get_network_boxes(net, custom_image_bgr.shape[1], custom_image_bgr.shape[0], thresh, hier_thresh, None, 0, pnum, letter_box) # OpenCV
        dets = get_network_boxes(net, im.w, im.h, thresh, hier_thresh, None, 0, pnum, letter_box)
        if debug: print("Got dets")
        num = pnum[0]
        if debug: print("got zeroth index of pnum")
        if nms:
            do_nms_sort(dets, num, meta.classes, nms)
        if debug: print("did sort")
        res = []
        if debug: print("about to range")
        for j in range(num):
            if debug: print("Ranging on "+str(j)+" of "+str(num))
            if debug: print("Classes: "+str(meta), meta.classes, meta.names)
            for i in range(meta.classes):
                if debug: print("Class-ranging on "+str(i)+" of "+str(meta.classes)+"= "+str(dets[j].prob[i]))
                if dets[j].prob[i] > 0:
                    b = dets[j].bbox
                    if altNames is None:
                        nameTag = meta.names[i]
                    else:
                        nameTag = altNames[i]
                    if debug:
                        print("Got bbox", b)
                        print(nameTag)
                        print(dets[j].prob[i])
                        print((b.x, b.y, b.w, b.h))
                    res.append((nameTag, dets[j].prob[i], (b.x, b.y, b.w, b.h)))
        if debug: print("did range")
        res = sorted(res, key=lambda x: -x[1])
        if debug: print("did sort")
        free_detections(dets, num)
        if debug: print("freed detections")
        return res


    netMain = None
    metaMain = None
    altNames = None

    def performDetect(imagePath=os.path.join(cwd,"dog.jpg"), thresh= 0.25, configPath = os.path.join(cwd,"yolov4.cfg"), weightPath = "yolov4.weights", metaPath= "coco.data", showImage= True, makeImageOnly = False, initOnly= False):
        """
        Convenience function to handle the detection and returns of objects.

        Displaying bounding boxes requires libraries scikit-image and numpy

        Parameters
        ----------------
        imagePath: str
            Path to the image to evaluate. Raises ValueError if not found

        thresh: float (default= 0.25)
            The detection threshold

        configPath: str
            Path to the configuration file. Raises ValueError if not found

        weightPath: str
            Path to the weights file. Raises ValueError if not found

        metaPath: str
            Path to the data file. Raises ValueError if not found

        showImage: bool (default= True)
            Compute (and show) bounding boxes. Changes return.

        makeImageOnly: bool (default= False)
            If showImage is True, this won't actually *show* the image, but will create the array and return it.

        initOnly: bool (default= False)
            Only initialize globals. Don't actually run a prediction.

        Returns
        ----------------------


        When showImage is False, list of tuples like
            ('obj_label', confidence, (bounding_box_x_px, bounding_box_y_px, bounding_box_width_px, bounding_box_height_px))
            The X and Y coordinates are from the center of the bounding box. Subtract half the width or height to get the lower corner.

        Otherwise, a dict with
            {
                "detections": as above
                "image": a numpy array representing an image, compatible with scikit-image
                "caption": an image caption
            }
        """
        # Import the global variables. This lets us instance Darknet once, then just call performDetect() again without instancing again
        global metaMain, netMain, altNames #pylint: disable=W0603
        assert 0 < thresh < 1, "Threshold should be a float between zero and one (non-inclusive)"
        if not os.path.exists(configPath):
            raise ValueError("Invalid config path `"+os.path.abspath(configPath)+"`")
        if not os.path.exists(weightPath):
            raise ValueError("Invalid weight path `"+os.path.abspath(weightPath)+"`")
        if not os.path.exists(metaPath):
            raise ValueError("Invalid data file path `"+os.path.abspath(metaPath)+"`")
        if netMain is None:
            netMain = load_net_custom(configPath.encode("ascii"), weightPath.encode("ascii"), 0, 1)  # batch size = 1
        if metaMain is None:
            metaMain = load_meta(metaPath.encode("ascii"))
        if altNames is None:
            # In Python 3, the metafile default access craps out on Windows (but not Linux)
            # Read the names file and create a list to feed to detect
            try:
                with open(metaPath) as metaFH:
                    metaContents = metaFH.read()
                    import re
                    match = re.search("names *= *(.*)$", metaContents, re.IGNORECASE | re.MULTILINE)
                    if match:
                        result = match.group(1)
                    else:
                        result = None
                    try:
                        if os.path.exists(result):
                            with open(result) as namesFH:
                                namesList = namesFH.read().strip().split("\n")
                                altNames = [x.strip() for x in namesList]
                    except TypeError:
                        pass
            except Exception:
                pass
        if initOnly:
            print("Initialized detector")
            return None
        if not os.path.exists(imagePath):
            raise ValueError("Invalid image path `"+os.path.abspath(imagePath)+"`")
        # Do the detection
        #detections = detect(netMain, metaMain, imagePath, thresh)	# if is used cv2.imread(image)
        detections = detect(netMain, metaMain, imagePath.encode("ascii"), thresh)
        if showImage:
            try:
                from skimage import io, draw
                import numpy as np
                image = io.imread(imagePath)
                print("*** "+str(len(detections))+" Results, color coded by confidence ***")
                imcaption = []
                for detection in detections:
                    label = detection[0]
                    confidence = detection[1]
                    pstring = label+": "+str(np.rint(100 * confidence))+"%"
                    imcaption.append(pstring)
                    print(pstring)
                    bounds = detection[2]
                    shape = image.shape
                    # x = shape[1]
                    # xExtent = int(x * bounds[2] / 100)
                    # y = shape[0]
                    # yExtent = int(y * bounds[3] / 100)
                    yExtent = int(bounds[3])
                    xEntent = int(bounds[2])
                    # Coordinates are around the center
                    xCoord = int(bounds[0] - bounds[2]/2)
                    yCoord = int(bounds[1] - bounds[3]/2)
                    boundingBox = [
                        [xCoord, yCoord],
                        [xCoord, yCoord + yExtent],
                        [xCoord + xEntent, yCoord + yExtent],
                        [xCoord + xEntent, yCoord]
                    ]
                    # Wiggle it around to make a 3px border
                    rr, cc = draw.polygon_perimeter([x[1] for x in boundingBox], [x[0] for x in boundingBox], shape= shape)
                    rr2, cc2 = draw.polygon_perimeter([x[1] + 1 for x in boundingBox], [x[0] for x in boundingBox], shape= shape)
                    rr3, cc3 = draw.polygon_perimeter([x[1] - 1 for x in boundingBox], [x[0] for x in boundingBox], shape= shape)
                    rr4, cc4 = draw.polygon_perimeter([x[1] for x in boundingBox], [x[0] + 1 for x in boundingBox], shape= shape)
                    rr5, cc5 = draw.polygon_perimeter([x[1] for x in boundingBox], [x[0] - 1 for x in boundingBox], shape= shape)
                    boxColor = (int(255 * (1 - (confidence ** 2))), int(255 * (confidence ** 2)), 0)
                    draw.set_color(image, (rr, cc), boxColor, alpha= 0.8)
                    draw.set_color(image, (rr2, cc2), boxColor, alpha= 0.8)
                    draw.set_color(image, (rr3, cc3), boxColor, alpha= 0.8)
                    draw.set_color(image, (rr4, cc4), boxColor, alpha= 0.8)
                    draw.set_color(image, (rr5, cc5), boxColor, alpha= 0.8)
                if not makeImageOnly:
                    io.imshow(image)
                    io.show()
                detections = {
                    "detections": detections,
                    "image": image,
                    "caption": "\n<br/>".join(imcaption)
                }
            except Exception as e:
                print("Unable to show image: "+str(e))
        return detections

    def performBatchDetect(thresh= 0.25, configPath = "./cfg/yolov3.cfg", weightPath = "yolov3.weights", metaPath= "./cfg/coco.data", hier_thresh=.5, nms=.45, batch_size=3):
        import cv2
        import numpy as np
        # NB! Image sizes should be the same
        # You can change the images, yet, be sure that they have the same width and height
        img_samples = ['data/person.jpg', 'data/person.jpg', 'data/person.jpg']
        image_list = [cv2.imread(k) for k in img_samples]

        net = load_net_custom(configPath.encode('utf-8'), weightPath.encode('utf-8'), 0, batch_size)
        meta = load_meta(metaPath.encode('utf-8'))
        pred_height, pred_width, c = image_list[0].shape
        net_width, net_height = (network_width(net), network_height(net))
        img_list = []
        for custom_image_bgr in image_list:
            custom_image = cv2.cvtColor(custom_image_bgr, cv2.COLOR_BGR2RGB)
            custom_image = cv2.resize(
                custom_image, (net_width, net_height), interpolation=cv2.INTER_NEAREST)
            custom_image = custom_image.transpose(2, 0, 1)
            img_list.append(custom_image)

        arr = np.concatenate(img_list, axis=0)
        arr = np.ascontiguousarray(arr.flat, dtype=np.float32) / 255.0
        data = arr.ctypes.data_as(POINTER(c_float))
        im = IMAGE(net_width, net_height, c, data)

        batch_dets = network_predict_batch(net, im, batch_size, pred_width,
                                                    pred_height, thresh, hier_thresh, None, 0, 0)
        batch_boxes = []
        batch_scores = []
        batch_classes = []
        for b in range(batch_size):
            num = batch_dets[b].num
            dets = batch_dets[b].dets
            if nms:
                do_nms_obj(dets, num, meta.classes, nms)
            boxes = []
            scores = []
            classes = []
            for i in range(num):
                det = dets[i]
                score = -1
                label = None
                for c in range(det.classes):
                    p = det.prob[c]
                    if p > score:
                        score = p
                        label = c
                if score > thresh:
                    box = det.bbox
                    left, top, right, bottom = map(int,(box.x - box.w / 2, box.y - box.h / 2,
                                                box.x + box.w / 2, box.y + box.h / 2))
                    boxes.append((top, left, bottom, right))
                    scores.append(score)
                    classes.append(label)
                    boxColor = (int(255 * (1 - (score ** 2))), int(255 * (score ** 2)), 0)
                    cv2.rectangle(image_list[b], (left, top),
                              (right, bottom), boxColor, 2)
            cv2.imwrite(os.path.basename(img_samples[b]),image_list[b])

            batch_boxes.append(boxes)
            batch_scores.append(scores)
            batch_classes.append(classes)
        free_batch_detections(batch_dets, batch_size)
        return batch_boxes, batch_scores, batch_classes    

    if __name__ == "__main__":
        print(performDetect())
        #Uncomment the following line to see batch inference working 
        #print(performBatchDetect())
