#!/usr/bin/python
# -*- coding: utf-8 -*-

import sys
import openface
import subprocess
import cv2
import uuid
import os
import random
import json
import pickle
import pdb
import dlib
import codecs
import time
import urllib
from ConfigParser import SafeConfigParser, NoOptionError

import numpy as np
np.set_printoptions(precision=2)


class Face:
    def __init__(self, box, rep):
        self.box = box
        self.rep = rep


class Person:
    def __init__(self, name, face, confidence):
        self.name = name
        self.face = face
        self.confidence = confidence


# Return the percentage two squares intersect each other with respect to the smallest box
# from http://stackoverflow.com/questions/9324339/how-much-do-two-rectangles-overlap
def squares_intersect(s1, s2):
    s1x1 = s1.left()
    s1y1 = s1.top()
    s1x2 = s1x1 + s1.height()
    s1y2 = s1y1 + s1.width()
    s2x1 = s2.left()
    s2y1 = s2.top()
    s2x2 = s2x1 + s2.height()
    s2y2 = s2y1 + s2.width()

    intersect = max(0, min(s1x2, s2x2) - max(s1x1, s2x1)) * max(0, min(s1y2, s2y2) - max(s1y1, s2y1))

    s1a = s1.height() * s1.width()
    s2a = s2.height() * s2.width()

    return float(intersect) / float(min(s1a, s2a))


# Use dlib to remove false positives from face boxes generated by opencv
def is_false_positive(img, box):
    box = cv2_rect_to_dlib(box)
    shape = np.shape(img)
    x1 = box.left()
    y1 = box.top()
    x2 = min(x1 + box.width(), shape[1])
    y2 = min(y1 + box.height(), shape[0])

    faces = align.getAllFaceBoundingBoxes(img)

    return len(faces) == 0


def cv2_rect_to_dlib(rect):
    x1 = long(rect[0])
    y1 = long(rect[1])
    x2 = long(rect[2]) + x1
    y2 = long(rect[3]) + y1
    
    return dlib.rectangle(x1, y1, x2, y2)


# Face detection using opencv. Give some false positives, but is very fast.
def get_faces_bounding_boxes_cv(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    face_boxes = face_cascade.detectMultiScale(
        gray, 
        cv_face_box_scale_factor, 
        cv_face_box_min_neighbours, 
        cv2.cv.CV_HAAR_DO_CANNY_PRUNING, 
        minSize=cv_face_box_min_size
    )

    boxes = []

    for box in face_boxes:
        if not is_false_positive(img, box):
            boxes.append(cv2_rect_to_dlib(box))
            print("Found face %s" % (box))
        else:
            print("Found false positive %s" % (box))

    return boxes


# Face detection using dlib. Accurate but slow.
def get_faces_bounding_boxes_dlib(img):
    # Convert from dlib.rectangles to list
    return [x for x in align.getAllFaceBoundingBoxes(img)]


# Returns a user with its face box intersecting box above the given threshold 
def get_tracked_person(box):
    for person in tracked_persons:
        if squares_intersect(person.face.box, box) > face_intersect_threshold:
            #print("%s and %s intersect with %s" % (person.face.box, box,squares_intersect(person.face.box, box)))
            return person

    return None
 

def get_faces(boxes, img):
    faces = []

    for box in boxes:
        tracked_person = get_tracked_person(box)

        # Only do face detection on faces which aren't already tracked
        if not tracked_person:
            aligned_face = align.align(face_image_dim, img, box, landmarkIndices=openface.AlignDlib.OUTER_EYES_AND_NOSE)
            rep = net.forward(aligned_face)
            faces.append(Face(box, rep))
        else:
            # Update the face box on the tracked face
            tracked_person.face.box = box

    return faces


# Play a welcome message
def optionally_play_message(person):
    if text_to_speech_function is None:
        return

    if not person.name in played_welcome_messages:
        played_welcome_messages[person.name] = 0

    if played_welcome_messages[person.name] + welcome_message_sleep_time < time.time():
        message = random.choice(available_welcome_messages[language]).replace("{name}", "%s") % (person.name)
        text_to_speech_function(message)
        played_welcome_messages[person.name] = time.time()

def save_unknown_face_img(img, face):
    if not os.path.exists('./generated/unknown'):
        os.makedirs('./generated/unknown')
                
    aligned_face = align.align(face_image_dim, img, face.box, landmarkIndices=openface.AlignDlib.OUTER_EYES_AND_NOSE)
    unknown_file = './generated/unknown/%s-%s.png' % (session_id, generated_image_id)
    cv2.imwrite(unknown_file, aligned_face)
    print("Saved unknown image %s" % (unknown_file))

def find_persons(faces, labels, classifier, img):
    global generated_image_id
    persons = []
    confidences = []
    
    for i, face in enumerate(faces):
        try:
            rep = face.rep.reshape(1, -1)
        except:
            # No Face detected
            return (None, None, None)

        predictions = classifier.predict_proba(rep).ravel()
        max_i = np.argmax(predictions)
        
        name = labels.inverse_transform(max_i)
        confidence = predictions[max_i]

        if confidence > person_confidence_threshold:
            person = Person(name, face, confidence)
            persons.append(person)

            if name != "-1":
                print("Added %s with confidence %s" % (name, confidence))
                optionally_play_message(person)
            else:
                print("Added unknown person with confidence %s" % (confidence))
                
        else:
            print("Ignored %s with confidence %s" % (name, confidence))

            if save_unknown_faces:
                save_unknown_face_img(img, face)

                generated_image_id += 1
        
    return persons


def draw_person_box(img, person):
    box = person.face.box
    x1 = int(box.left())
    y1 = int(box.top())
    x2 = int(box.right())
    y2 = int(box.bottom())
    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)

    ts, _ = cv2.getTextSize(person.name, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 1)
    cv2.putText(img, person.name, (x1 - ts[0]/2 + (x2-x1)/2, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,0,0), 1)


# Match persons and face boxes, remove persons with no corresponding face box
def prune_match_boxes_persons(boxes, persons):
    pruned_boxes = boxes[:]
    pruned_persons = []

    for person in persons:
        for i, box in enumerate(boxes):
            intersect = squares_intersect(person.face.box, box)
            if intersect > face_intersect_threshold:
                person.face.box = box
                pruned_persons.append(person)
                pruned_boxes[i] = None

    pruned_boxes = [box for box in pruned_boxes if box is not None]

    return (pruned_boxes, pruned_persons)


def execute_next_command():
    global command_queue, running_command

    if running_command:
        status = running_command.poll()
        if status is not None:
            running_command = None

    elif len(command_queue) > 0:
        command = command_queue[0]
        command_queue = command_queue[1:]
        
        process = subprocess.Popen(command, shell=True)
        running_command = process


def schedule_command(command):
    global command_queue

    command_queue.append(command)
        

def marytts_speech(text):
    command = 'curl "http://localhost:59125/process?INPUT_TYPE=TEXT&AUDIO=WAVE_FILE&OUTPUT_TYPE=AUDIO&LOCALE=%s&INPUT_TEXT=%22%s%22"|aplay' % (language, urllib.quote_plus(text))
    schedule_command(command)


def espeak_speech(text):
    locale_str = ""

    if language == 'sv':
        locale_str = "-vsv "

    command = 'espeak %s"%s"' % (locale_str, text)
    schedule_command(command.encode('utf-8'))


if __name__ == '__main__':
    conf_file = 'default.conf'

    if len(sys.argv) > 1:
        conf_file = sys.argv[1]

    config = SafeConfigParser()
    config.readfp(codecs.open(conf_file, "r"))

    face_detector_conf = config.get('FaceDetection', 'detector')
    
    if face_detector_conf == "dlib":    
        face_detector = get_faces_bounding_boxes_dlib
    elif face_detector_conf == "opencv":
        face_detector = get_faces_bounding_boxes_cv

    face_intersect_threshold = config.getfloat('FaceDetection', 'intersect_threshold')
    face_cascade = cv2.CascadeClassifier(config.get('FaceDetection', 'cascade'))
    cv_face_box_min_size = tuple(map(int, config.get('FaceDetection', 'min_box_size').split(',')))
    cv_face_box_scale_factor = config.getfloat('FaceDetection', 'scale_factor')
    cv_face_box_min_neighbours = config.getfloat('FaceDetection', 'min_neighbours')

    face_predictor_file = config.get('Embedding', 'predictor')
    torch_network_model_file = config.get('Embedding', 'model')
    face_image_dim = config.getint('Embedding', 'image_dim')

    show_video = config.getboolean('Video', 'show_video')
    rotate_video = config.getint('Video', 'rotate_video')
    video_capture_device = config.getint('Video', 'device')
    image_size = tuple(map(int, config.get('Video', 'image_size').split(',')))
    save_unknown_faces = config.getboolean('Video', 'save_unknown')

    person_confidence_threshold = config.getfloat('Identification', 'min_confidence')
    classifierFile = config.get('Identification', 'classifier')

    language = config.get('Greetings', 'language')
    welcome_message_sleep_time = config.getint('Greetings', 'message_wait_time')
    available_welcome_messages = json.loads(config.get('Greetings', 'messages').replace('\n', ''))

    try:
        speech_fun_name = config.get('Greetings', 'speech_api')
    except NoOptionError:
        speech_fun_name = None

    if speech_fun_name == 'espeak':
        text_to_speech_function = espeak_speech
    elif speech_fun_name == 'marytts':
        text_to_speech_function = marytts_speech
    else:
        text_to_speech_function = None

    update_faces_skip_frames = config.getint('Performance', 'skip_frames')

    command_queue = []
    running_command = None
    played_welcome_messages = {}
    generated_image_id = 0
    align = openface.AlignDlib(face_predictor_file)    
    session_id = str(uuid.uuid1())
    tracked_persons = []
    net = openface.TorchNeuralNet(torch_network_model_file, imgDim=face_image_dim)
    iteration = 0
     
    vc = cv2.VideoCapture(video_capture_device)

    with open(classifierFile, 'r') as f:
        (labels, classifier) = pickle.load(f)

    while True:
        _, img = vc.read()

        # Rotate the screen if the camera is tilted
        cols, rows, _ = img.shape

        #pdb.set_trace()
        if rotate_video > 0:
            M = cv2.getRotationMatrix2D((cols/2, rows/2),rotate_video,1)
            img = cv2.warpAffine(img, M, (cols, rows))
        
        img = cv2.resize(img, image_size, interpolation = cv2.INTER_CUBIC)

        # Skip frames to avoid these expensive steps
        if iteration % update_faces_skip_frames == 0:
            boxes = face_detector(img)
            boxes, pruned_tracked_persons = prune_match_boxes_persons(boxes, tracked_persons)
 
            faces = get_faces(boxes, img)
            tracked_persons = find_persons(faces, labels, classifier, img) + pruned_tracked_persons

        for person in tracked_persons:
            if show_video:
                draw_person_box(img, person)
            
            print("Tracking %s in %s" % (person.name, person.face.box))

        if show_video:
            cv2.imshow('Video', img)

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        execute_next_command()

        iteration += 1

    vc.release()
    cv2.destroyAllWindows()
