"""
Detection.py
"""

## Libraries
from ctypes import *
import random
import cv2
# from PIL import Image
# import datetime
import os
# import tensorflow as tf
from pathlib2 import Path
from gittle import Gittle

# Select GPU devices to activate them and set maximum usage of the memory
os.environ["CUDA_VISIBLE_DEVICES"] = '3,2'


# config = tf.ConfigProto()
# config.gpu_options.per_process_gpu_memory_fraction = 0.9
# session = tf.Session(config=config)


## Classes


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


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


ROOTPATH, filename = os.path.split(os.path.abspath(__file__))
SoPath = os.path.join(ROOTPATH, 'libdarknet.so')

so_file = Path(SoPath)
if so_file.is_file() == False:
    print(False)
    repo_url = 'http://iip.whu.edu.cn/GitLab/Huairui/darknet_cfg.git'
    Gittle.clone(repo_url, ROOTPATH)
lib = CDLL(SoPath, RTLD_GLOBAL)
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

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

set_gpu = lib.cuda_set_device
set_gpu.argtypes = [c_int]

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)]
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_ptrs = lib.free_ptrs
free_ptrs.argtypes = [POINTER(c_void_p), c_int]

network_predict = lib.network_predict
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

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)


def GenerateNet():
    ROOTPATH, filename = os.path.split(os.path.abspath(__file__))
    SoPath = os.path.join(ROOTPATH, 'libdarknet.so')
    so_file = Path(SoPath)
    if so_file.is_file() == False:
        repo_url = 'http://iip.whu.edu.cn/GitLab/Huairui/darknet_cfg.git'
        Gittle.clone(repo_url, ROOTPATH)
        GenerateData()
    cfg_path = os.path.join(ROOTPATH, 'cfg/yolov3.cfg')
    weight_path = os.path.join(ROOTPATH, 'backup/yolov3.weights')
    meta_path = os.path.join(ROOTPATH, 'cfg/coco2.data')
    pic_path = os.path.join(ROOTPATH, 'data/kite.jpg')
    net = load_net(cfg_path.encode('UTF-8'), weight_path.encode('UTF-8'), 0)
    meta = load_meta(meta_path.encode('UTF-8'))


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


def detect(net, meta, image, thresh=.5, hier_thresh=.5, nms=.45):
    im = load_image(image, 0, 0)
    num = c_int(0)
    pnum = pointer(num)
    predict_image(net, im)
    dets = get_network_boxes(net, im.w, im.h, thresh, hier_thresh, None, 0, pnum)
    num = pnum[0]
    if (nms): do_nms_obj(dets, num, meta.classes, nms);

    res = []
    for j in range(num):
        for i in range(meta.classes):
            if dets[j].prob[i] > 0:
                b = dets[j].bbox
                res.append((meta.names[i], dets[j].prob[i], (b.x, b.y, b.w, b.h)))
    res = sorted(res, key=lambda x: -x[1])
    free_image(im)
    free_detections(dets, num)
    return res


# display the pic after detecting
def GenerateResult(file_name, r, VW, flag):
    img = cv2.imread(file_name)

    for i in range(len(r)):
        x1 = r[i][2][0] - r[i][2][2] / 2
        y1 = r[i][2][1] - r[i][2][3] / 2
        x2 = r[i][2][0] + r[i][2][2] / 2
        y2 = r[i][2][1] + r[i][2][3] / 2

        cv2.rectangle(img, (int(x1 - 1), int(y1 - 21)), (int(x1 + 70), int(y1)), (0, 255, 255), -1)
        cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 255), 2)
        cv2.putText(img, str(r[i][0]), (int(x1 + 5), int(y1 - 4)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)

    # if flag == 1:
    #     cv2.imwrite("../video/frame/original.jpg", img)
    #     flag += 1
    VW.write(img)


def ProcessVideo(net, meta, video):
    capture = cv2.VideoCapture(video)
    if capture.isOpened():
        success = True
    else:
        success = False

    frame_count = 1

    print("Pocessing frames...")
    fps = 25
    fourcc = cv2.VideoWriter_fourcc(*'MJPG')
    VideoWriter = cv2.VideoWriter('../video/processed.avi', fourcc, fps, resolution)
    flag = 1
    while (success and frame_count <= 125):
        print("Processing frame %d..." % frame_count)
        success, frame = capture.read()

        frame_name = "../video/cache/%d.jpg" % frame_count

        img = Image.fromarray(frame)
        img.save(frame_name)
        begin = datetime.datetime.now()
        r = detect(net, meta, frame_name)
        end = datetime.datetime.now()
        # ProcessingTime.append((end - begin).total_seconds())
        GenerateResult(frame_name, r, VideoWriter, flag)

        frame_count += 1

    VideoWriter.release()
    capture.release()
    print("task finished.")


def ProcessImage(img_path, net, meta):
    img_path_d = img_path.encode('UTF-8')
    r = detect(net, meta, img_path_d)
    img = cv2.imread(img_path)
    for i in range(len(r)):
        x1 = r[i][2][0] - r[i][2][2] / 2
        y1 = r[i][2][1] - r[i][2][3] / 2
        x2 = r[i][2][0] + r[i][2][2] / 2
        y2 = r[i][2][1] + r[i][2][3] / 2

        cv2.rectangle(img, (int(x1 - 1), int(y1 - 21)), (int(x1 + 70), int(y1)), (0, 255, 255), -1)
        cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 255), 2)
        cv2.putText(img, r[i][0].decode('UTF-8'), (int(x1 + 5), int(y1 - 4)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0),
                    2)
    # cv2.imshow('Figure Window', img)
    # cv2.waitKey(0)
    cv2.imwrite('result.jpg', img)

    # cv2.destroyAllWindows()


def GenerateData():
    dataPath = os.path.join(ROOTPATH, 'cfg/coco2.data')
    namePath = os.path.join(ROOTPATH, 'data/coco.names')
    with open(dataPath, 'wb') as f:
        f.write('classes= 80\n'.encode('UTF-8'))
        f.write('train = /home/pjreddie/data/coco/trainvalno5k.txt\n'.encode('UTF-8'))
        f.write('valid= coco_testdev\n'.encode('UTF-8'))
        f.write('names= ' + namePath.encode('UTF-8') + '\n'.encode('UTF-8'))
        f.write('backup= '.encode('UTF-8') + ROOTPATH.encode('UTF-8'))


def DeleteCacheFiles():
    directory = "../video/cache"
    os.chdir(directory)
    cwd = os.getcwd()
    files = os.listdir(cwd)
    for file in files:
        os.remove(file)
    print("Cache clear succeeded.")


def GenerateNetwork():
    net = load_net(cfg_path, weight_path, 0)
    meta = load_meta(meta_path)
    return net, meta


## Body

# resolution = (1280, 720)
GenerateData()
cfg_path = os.path.join(ROOTPATH, 'cfg/yolov3.cfg')
weight_path = os.path.join(ROOTPATH, 'backup/yolov3.weights')
meta_path = os.path.join(ROOTPATH, 'cfg/coco2.data')
pic_path = os.path.join(ROOTPATH, 'data/kite.jpg')

# net = load_net(cfg_path, weight_path, 0)
# meta = load_meta(meta_path)
# ProcessImage(pic_path)
