"""
Detect objects and faces or persons using tensorflow-tpu served by zerorpc.

This needs to be called from a zerorpc client with
an array of alarm frame image paths.

This is part of the smart-zoneminder project.
See https://github.com/goruck/smart-zoneminder

Copyright (c) 2018 ~ 2020 <NAME>
"""

import numpy as np
import json
import zerorpc
import pickle
import cv2
import logging
import gevent
import face_recognition
import sys
import tflite_runtime.interpreter as tflite
from signal import SIGINT, SIGTERM
from edgetpu.detection.engine import DetectionEngine

logging.basicConfig(level=logging.INFO)

# Open configuration file. 
with open('./config.json') as fp:
    config = json.load(fp)

obj_config = config['objDetServer']
face_config = config['faceDetServer']
person_config = config['personClassServer']

### Object detection configuration. ###
# Tensorflow object and face detection file system paths.
OBJ_MODEL = obj_config['objModelPath']
OBJ_LABEL_MAP = obj_config['labelMapPath']
# If consecutive ZoneMinder image frames are found then skip this many after the first.
OBJ_CON_IMG_SKIP = obj_config['conseqImagesToSkip']
# Minimum score for valid TF object detection. 
OBJ_MIN_SCORE_THRESH = obj_config['minScore']
# IPC (or TCP) socket for zerorpc.
# This must match the zerorpc client config.
OBJ_ZRPC_PIPE = obj_config['zerorpcPipe']

### Face detection configuration. ###
# Tensorflow face detection model path.
FACE_DET_MODEL = face_config['faceDetModelPath']
# py torch face embeddings model path (if used).
FACE_EMB_MODEL = face_config['faceEmbModelPath']
# IPC (or TCP) socket for zerorpc.
# This must match the zerorpc client config.
FACE_ZRPC_PIPE = face_config['zerorpcPipe']
# Settings for face classifier.
# The model and label encoder needs to be generated by 'train.py' first. 
FACE_CLASS_MODEL = face_config['modelPath']
FACE_LABEL_MAP = face_config['labelPath']
FACE_MIN_PROBA = face_config['minProba']
# Images with Variance of Laplacian less than this are declared blurry. 
FACE_FOCUS_MEASURE_THRESHOLD = face_config['focusMeasureThreshold']
# Faces with width or height less than this are too small for recognition.
# In pixels.
FACE_MIN = face_config['minFace']
# Number of times to resample for dlib face encoder.
FACE_NUM_JITTERS = face_config['numJitters']

### Person classifier configuration. ###
# TFLite model path. 
PERSON_CLASS_MODEL = person_config['personClassModelPath']
# Label map path.
PERSON_LABEL_MAP = person_config['labelMap']
# Classification threshold.
PERSON_MIN_PROBA = person_config['minProba']
# IPC (or TCP) socket for zerorpc.
# This must match the zerorpc client config.
PERSON_ZRPC_PIPE = person_config['zerorpcPipe']

### Global configuration parameters.
# Use dlib-based face detector or cnn-based person detector.
RECOGNIZE_MODE = config['recognizeMode']
# Mount point of zm alarms on local tpu machine. 
MOUNT_POINT = config['mountPoint']
# Heartbeat interval for zerorpc client in ms.
# This must match the zerorpc client config. 
ZRPC_HEARTBEAT = config['zerorpcHeartBeat']

def ReadLabelFile(file_path):
    # Function to read labels from text files.
    with open(file_path, 'r') as f:
        lines = f.readlines()
    ret = {}
    for line in lines:
        pair = line.strip().split(maxsplit=1)
        ret[int(pair[0])] = pair[1].strip()
    return ret

def face_classifier(recognizer, le, encoding, min_proba):
    # perform classification to recognize the face based on 128D encoding
    # note: reshape(1,-1) converts 1D array into 2D
    preds = recognizer.predict_proba(encoding.reshape(1, -1))[0]
    j = np.argmax(preds)
    proba = preds[j]
    logging.debug('face classifier proba {} name {}'.format(proba, le.classes_[j]))
    if proba >= min_proba:
        name = le.classes_[j]
        logging.debug('face classifier says this is {}'.format(name))
    else:
        name = None # prob too low to recog face
        logging.debug('face classifier cannot recognize face')
    return name, proba

def skip_inference(frame_num, monitor, labels, image_path, objects_in_image):
    """
    If consecutive frames then repeat last label and skip a new inference.
    
    Image paths must be in the form of:
    '/nvr/zoneminder/events/BackPorch/18/06/20/19/20/04/00224-capture.jpg'.
    """

    old_frame_num = frame_num
    old_monitor = monitor
    skip = False

    if OBJ_CON_IMG_SKIP == 0: return skip, frame_num, monitor

    try:
        frame_num = int((image_path.split('/')[-1]).split('-')[0])
        monitor = image_path.split('/')[4]
    except (ValueError, IndexError):
        logging.error('Could not derive information from image path.')
        objects_in_image.append({'image': image_path, 'labels': []})
        skip = True
        return skip, frame_num, monitor
                    
    # Only apply skip logic if frames are from the same monitor. 
    if monitor == old_monitor:
        # Only apply skip logic if alarm frames are from the same event.
        # Intra-event frames are monotonically increasing.
        frame_diff = frame_num - old_frame_num
        if frame_diff > 0:
            # Skip CON_IMG_SKIP frames after the first one. 
            if frame_diff <= OBJ_CON_IMG_SKIP:
                objects_in_image.append({'image': image_path, 'labels': labels})
                logging.debug('monitor {} old_monitor {} frame_num {} old_frame_num {}'
                    .format(monitor,old_monitor,frame_num,old_frame_num))
                logging.debug('Consecutive frame {}, skipping detect and copying previous labels.'
                    .format(frame_num))
                skip = True
                        
    return skip, frame_num, monitor

def resize_to_square(img, size, keep_aspect_ratio=False, interpolation=cv2.INTER_AREA):
    # Resize image to square shape.
    # If keep_aspect_ratio=True, then:
    #   If the original image is lanscape, add black pixels on the bottom-side only.
    #   If the original image is portrait, add black pixels on the right-side only.
    (h, w) = img.shape[:2]

    if h == w or keep_aspect_ratio == False:
        return cv2.resize(img, (size, size), interpolation)

    # Check if image is color. 
    chan = None if len(img.shape) < 3 else img.shape[2]

    # Determine size of black mask.
    mask_size = h if h > w else w

    if chan is None:
        mask = np.zeros((mask_size, mask_size), dtype=img.dtype)
        mask[:h, :w] = img[:h, :w]
    else:
        mask = np.zeros((mask_size, mask_size, chan), dtype=img.dtype)
        mask[:h, :w, :] = img[:h, :w, :]

    return cv2.resize(mask, (size, size), interpolation)

# zerorpc obj det server.
class ObjDetectRPC(object):
    def __init__(self):
        self.obj_engine = DetectionEngine(OBJ_MODEL)
        self.labels_map = ReadLabelFile(OBJ_LABEL_MAP)

    def detect_objects(self, test_image_paths):
        objects_in_image = [] # holds all objects found in image
        labels = [] # labels of detected objects
        frame_num = 0 # ZoneMinder current alarm frame number
        monitor = '' # ZoneMinder current monitor name

        for image_path in test_image_paths:
            logging.debug('**********Find object(s) for {}'.format(image_path))

            # If consecutive frames then repeat last label and skip inference.
            # This behavior controlled by CON_IMG_SKIP.
            skip, frame_num, monitor = skip_inference(frame_num, monitor,
                labels, image_path, objects_in_image)
            if skip is True:
                continue

            # Read image from disk. 
            img = cv2.imread(MOUNT_POINT + image_path)
            #cv2.imwrite('./obj_img.jpg', img)
            if img is None:
                # Bad image was read.
                logging.error('Bad image was read.')
                objects_in_image.append({'image': image_path, 'labels': []})
                continue

            # Resize. The tpu obj det requires (300, 300).
            res = resize_to_square(img=img, size=300, keep_aspect_ratio=True,
                interpolation=cv2.INTER_AREA)
            #cv2.imwrite('./obj_res.jpg', res)

            # Run object inference.
            detection = self.obj_engine.detect_with_input_tensor(res.reshape(-1),
                threshold=0.05, top_k=3)

            # Get labels and scores of detected objects.
            labels = [] # new detection, clear labels list. 
            (h, w) = img.shape[:2] # use original image size for box coords
            for obj in detection:
                logging.debug('id: {} name: {} score: {}'
                    .format(obj.label_id, self.labels_map[obj.label_id], obj.score))
                if obj.score > OBJ_MIN_SCORE_THRESH:
                    object_dict = {}
                    object_dict['id'] = obj.label_id
                    object_dict['name'] = self.labels_map[obj.label_id]
                    object_dict['score'] = float(obj.score)
                    (xmin, ymin, xmax, ymax) = (obj.bounding_box.flatten().tolist()) * np.array([w, h, w, h])
                    object_dict['box'] = {'ymin': ymin, 'xmin': xmin, 'ymax': ymax, 'xmax': xmax}
                    labels.append(object_dict)

            objects_in_image.append({'image': image_path, 'labels': labels})
        return json.dumps(objects_in_image)

# zerorpc face detection server.
class FaceDetectRPC(object):
    def __init__(self):
        # Init TPU engine.
        self.face_engine = DetectionEngine(FACE_DET_MODEL)

        # Load face recognition model and the label encoder.
        with open(FACE_CLASS_MODEL, 'rb') as fp:
            self.recognizer = pickle.load(fp)
        with open(FACE_LABEL_MAP, 'rb') as fp:
            self.le = pickle.load(fp)

    def detect_faces(self, test_image_paths):
        # List that will hold all images with any face detection information. 
        objects_detected_faces = []

        # Loop over the images paths provided. 
        for obj in test_image_paths:
            logging.debug('**********Find Face(s) for {}'.format(obj['image']))
            for label in obj['labels']:
                # If the object detected is a person then try to identify face. 
                if label['name'] == 'person':
                    # Read image from disk. 
                    img = cv2.imread(MOUNT_POINT + obj['image'])
                    if img is None:
                        # Bad image was read.
                        logging.error('Bad image was read.')
                        label['face'] = None
                        continue

                    # First bound the roi using the coord info passed in.
                    # The roi is area around person(s) detected in image.
                    # (x1, y1) are the top left roi coordinates.
                    # (x2, y2) are the bottom right roi coordinates.
                    y2 = int(label['box']['ymin'])
                    x1 = int(label['box']['xmin'])
                    y1 = int(label['box']['ymax'])
                    x2 = int(label['box']['xmax'])
                    roi = img[y2:y1, x1:x2, :]
                    #cv2.imwrite('./roi.jpg', roi)
                    if roi.size == 0:
                        # Bad object roi...move on to next image.
                        logging.error('Bad object roi.')
                        label['face'] = None
                        continue

                    # Need roi shape for later conversion of face coords.
                    (h, w) = roi.shape[:2]
                    # Resize roi for face detection.
                    # The tpu face det model used requires (320, 320).
                    res = resize_to_square(img=roi, size=320, keep_aspect_ratio=True,
                        interpolation=cv2.INTER_AREA)
                    #cv2.imwrite('./res.jpg', res)

                    # Detect the (x, y)-coordinates of the bounding boxes corresponding
                    # to a face in the input image using the TPU engine.
                    # Its assumed that only one face is in the image. 
                    # NB: reshape(-1) converts the np img array into 1-d. 
                    detection = self.face_engine.detect_with_input_tensor(res.reshape(-1),
                        threshold=0.05, top_k=1)
                    if not detection:
                        # No face detected...move on to next image.
                        logging.debug('No face detected.')
                        label['face'] = None
                        continue
                        
                    # Convert coords and carve out face roi.
                    box = (detection[0].bounding_box.flatten().tolist()) * np.array([w, h, w, h])
                    (face_left, face_top, face_right, face_bottom) = box.astype('int')
                    face_roi = roi[face_top:face_bottom, face_left:face_right, :]
                    #cv2.imwrite('./face_roi.jpg', face_roi)
                    (f_h, f_w) = face_roi.shape[:2]
                    # If face width or height are not sufficiently large then skip.
                    if f_h < FACE_MIN or f_w < FACE_MIN:
                        logging.debug('Face too small to recognize.')
                        label['face'] = None
                        continue

                    # Compute the focus measure of the face
                    # using the Variance of Laplacian method.
                    # See https://www.pyimagesearch.com/2015/09/07/blur-detection-with-opencv/
                    gray = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
                    fm = cv2.Laplacian(gray, cv2.CV_64F).var()
                    # If fm below a threshold then face probably isn't clear enough
                    # for face recognition to work, so skip it. 
                    if fm < FACE_FOCUS_MEASURE_THRESHOLD:
                        logging.debug('Face too blurry to recognize.')
                        label['face'] = None
                        continue

                    # Find the 128-dimension face encoding for face in image.
                    # Convert image roi from BGR (OpenCV ordering) to dlib ordering (RGB).
                    rgb = cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)
                    # Convert face bbox into dlib format.
                    boxes = [(face_top, face_right, face_bottom, face_left)]
                    # Generate encodings. Only one face is assumed so take the 1st element. 
                    encoding = face_recognition.face_encodings(face_image=rgb,
                        known_face_locations=boxes, num_jitters=FACE_NUM_JITTERS)[0]
                    logging.debug('face encoding {}'.format(encoding))
                    # Perform svm classification on the encodings to recognize the face.
                    (name, proba) = face_classifier(
                        recognizer=self.recognizer,
                        le=self.le,
                        encoding=encoding,
                        min_proba=FACE_MIN_PROBA)

                    # Add face name to label metadata.
                    label['face'] = name
                    # Add face confidence to label metadata.
                    # (First convert NumPy value to native Python type for json serialization.)
                    label['faceProba'] = proba.item()
            # Add processed image to output list. 
            objects_detected_faces.append(obj)
        # Convert json to string and return data. 
        return(json.dumps(objects_detected_faces))

# zerorpc person classifier server
class PersonClassRPC(object):
    def __init__(self):
        # Load TFLite model and allocate tensors.
        self.interpreter = tflite.Interpreter(model_path=PERSON_CLASS_MODEL,
            experimental_delegates=[tflite.load_delegate('libedgetpu.so.1')])
        self.interpreter.allocate_tensors()

        # Get input and output tensors.
        self.input_details = self.interpreter.get_input_details()
        self.output_details = self.interpreter.get_output_details()

    def detect_faces(self, test_image_paths):
        # List that will hold all images with any person classification information. 
        objects_classified_persons = []

        # Loop over the images paths provided. 
        for obj in test_image_paths:
            logging.debug('**********Classify person for {}'.format(obj['image']))
            for label in obj['labels']:
                # If the object detected is a person then try to identify face. 
                if label['name'] == 'person':
                    # Read image from disk. 
                    img = cv2.imread(MOUNT_POINT + obj['image'])
                    if img is None:
                        # Bad image was read.
                        logging.error('Bad image was read.')
                        label['face'] = None
                        continue

                    # First bound the roi using the coord info passed in.
                    # The roi is area around person(s) detected in image.
                    # (x1, y1) are the top left roi coordinates.
                    # (x2, y2) are the bottom right roi coordinates.
                    y2 = int(label['box']['ymin'])
                    x1 = int(label['box']['xmin'])
                    y1 = int(label['box']['ymax'])
                    x2 = int(label['box']['xmax'])
                    roi = img[y2:y1, x1:x2, :]
                    if roi.size == 0:
                        # Bad object roi...move on to next image.
                        logging.error('Bad object roi.')
                        label['face'] = None
                        continue

                    # Resize roi for person classification model.
                    input_shape = self.input_details[0]['shape']
                    input_size = tuple(input_shape[1:3])
                    roi = cv2.resize(roi, input_size) # uint8
                    #cv2.imwrite('./roi.jpg', roi)

                    # Expand dimensions. 
                    roi = np.expand_dims(roi, axis=0)

                    # Classify person using the TPU engine.
                    # Its assumed that only one person is in the image. 
                    self.interpreter.set_tensor(self.input_details[0]['index'], roi)
                    self.interpreter.invoke()
                    classification = self.interpreter.get_tensor(self.output_details[0]['index'])[0]
                    
                    # Find most likely prediction.
                    class_id = np.argmax(classification)
                    proba = classification[class_id] / 256. # probas range from 1 to 256
                    person = PERSON_LABEL_MAP[class_id]
                    logging.debug('person classifier proba {} name {}'.format(proba, person))
                    if proba >= PERSON_MIN_PROBA:
                        name = person
                        logging.debug('person classifier says this is {}'.format(name))
                    else:
                        name = None # prob too low to recog face
                        logging.debug('person classifier cannot recognize person')

                    # Add face name to label metadata.
                    label['face'] = name
                    # Add face confidence to label metadata.
                    # (First convert NumPy value to native Python type for json serialization.)
                    label['faceProba'] = proba.item()
            # Add processed image to output list. 
            objects_classified_persons.append(obj)
        # Convert json to string and return data. 
        return(json.dumps(objects_classified_persons))

# Setup face detection or person classifier server.
if RECOGNIZE_MODE == 'person':
    zerorpc_obj = PersonClassRPC()
    zerorpc_pipe = PERSON_ZRPC_PIPE
elif RECOGNIZE_MODE == 'face':
    zerorpc_obj = FaceDetectRPC()
    zerorpc_pipe = FACE_ZRPC_PIPE
else:
    logging.error('Unknown recognizer mode.')
    sys.exit()
face_s = zerorpc.Server(zerorpc_obj, heartbeat=ZRPC_HEARTBEAT)
face_s.bind(zerorpc_pipe)
# Register graceful ways to stop server. 
gevent.signal(SIGINT, face_s.stop) # Ctrl-C
gevent.signal(SIGTERM, face_s.stop) # termination

# Setup object detection server.
obj_s = zerorpc.Server(ObjDetectRPC(), heartbeat=ZRPC_HEARTBEAT)
obj_s.bind(OBJ_ZRPC_PIPE)
# Register graceful ways to stop server. 
gevent.signal(SIGINT, obj_s.stop) # Ctrl-C
gevent.signal(SIGTERM, obj_s.stop) # termination

# Startup both servers.
# This will block until a gevent SIGINT or SIGTERM signal is caught.
gevent.joinall([gevent.spawn(face_s.run), gevent.spawn(obj_s.run)])