#!/usr/bin/python
from __future__ import print_function

from PIL import Image

import tkinter, tkinter.simpledialog
from fsdk import FSDK
from tk import MainWindow, Style

DRAW_FEATURES = False
trackerMemoryFile = "tracker70.dat"

print("Initializing FSDK... ", end='')
FSDK.ActivateLibrary(""); 
FSDK.Initialize()
print("OK\nLicense info:", FSDK.GetLicenseInfo())

FSDK.InitializeCapturing()
print('Looking for video cameras... ', end = '')
camList = FSDK.ListCameraNames()

if not camList:
    print("Please attach a camera."); 
    quit(1)

print(camList[0].devicePath)

camera = camList[0] # choose the first camera (0)
print("using '%s'" % camera)
formatList = FSDK.ListVideoFormats(camera)
print(*formatList, sep='\n')

vfmt = formatList[0] # choose the first format: vfmt.Width, vfmt.Height, vfmt.BPP
print('Selected camera format:', vfmt)
FSDK.SetVideoFormat(camera, vfmt)

print("Trying to open '%s'... " % camera, end = '')
camera = FSDK.OpenVideoCamera(camera)
print("OK")

try:
    fsdkTracker = FSDK.Tracker.FromFile(trackerMemoryFile)
except:
    fsdkTracker = FSDK.Tracker()  # creating a FSDK Tracker 

fsdkTracker.SetParameters( # set realtime face detection parameters
    RecognizeFaces=True, DetectFacialFeatures=DRAW_FEATURES,
    InternalResizeWidth=256, FaceDetectionThreshold=5,
)
fsdkTracker.SetParameter('SmoothAttributeLiveness', True); # use smooth minimum function for liveness values
fsdkTracker.SetParameter('AttributeLivenessSmoothingAlpha', 1); # smooth minimum parameter, 0 -> mean, inf -> min
fsdkTracker.SetParameter('LivenessFramesCount', 15); # minimal number of frames required to output liveness attribute


need_to_exit = False

def onDestroy():
    global need_to_exit
    need_to_exit = True

root = tkinter.Tk()
root.title('Live Recognition')
root.protocol("WM_DELETE_WINDOW", onDestroy)
root.geometry(f"{vfmt.Width}x{vfmt.Height}")

window = MainWindow(root)

activeStyle = Style(color='#0000FF', width=3, fill='')
faceStyle = Style(color='#FFFFFF', width=4, fill='')
featureStyle = Style(color='#FFFF60', width=2, fill='#FFFF60')
textColor = Style(color='#FFFFFF', width=1, fill='')
textShadow = Style(color='#808080', width=1, fill='')
livenessStyle = Style(color='#FFFFFF', width=1, fill='#808080')
livenessActiveStyle = Style(color='#FFFFFF', width=1, fill='#b0b0b0')

livenessButton = (2, 2, 240, 40)

class FaceFrame:

    def __init__(self, xc, yc, width, height):
        self.xc = xc
        self.yc = yc
        self.width  = int(width)
        self.height = int(height)


    def scale(self, scale=None):
        try:
            scale_w, scale_h = scale
        except:
            scale_w = scale_h = scale

        return FaceFrame(self.xc, self.yc, self.width * scale_w, self.height * scale_h)


    # axes aligned bounding box
    def aabb(self):
        w = self.width // 2
        h = self.height // 2

        return (self.xc - w, self.yc - h, self.xc + w, self.yc + h)


    def contains(self, x, y):
        return 2 * abs(x - self.xc) <= self.width and 2 * abs(y - self.yc) <= self.height


class FaceLocator:

    SCALE_W = 1
    SCALE_H = 1.25

    def __init__(self, face_id):
        try:
            fsdkTracker.LockID(face_id)
            self.name = fsdkTracker.GetName(face_id) or None
            fsdkTracker.UnlockID(face_id)
        except FSDK.IdNotFound:
            self.name = None

        self.frame = None
        self.features = []
        self.faceRect = None
        self.activeRect = None
        self.faceName = None
        self.face_id = face_id


    def draw(self):
        window.deleteObject(self.faceRect)
        window.deleteObject(self.activeRect)
        window.deleteObject(self.faceName)

        try:
            face_pos = fsdkTracker.GetFacePosition(0, self.face_id)
            self.frame = FaceFrame(face_pos.xc, face_pos.yc, face_pos.w, face_pos.w).scale(scale=(self.SCALE_W, self.SCALE_H))

            aabb = self.frame.aabb()
            self.faceRect = window.drawRectangle(*aabb, style=faceStyle)

            if mouse_pos is not None and self.frame.contains(*mouse_pos):
                self.activeRect = window.drawRectangle(*aabb, style=activeStyle)

            if self.name is not None:
                name = self.name
                if liveness_status:
                    try:
                        liveness = fsdkTracker.GetTrackerFacialAttribute(0, self.face_id, 'Liveness')
                        liveness = FSDK.GetValueConfidence(liveness, 'Liveness')
                        name = f"{name}\n({liveness * 100:.2f}% liveness)"

                    except FSDK.AttributeNotDetected:
                        pass
                self.faceName = window.drawText((self.frame.xc, self.frame.yc + self.frame.height // 2 + 25), name, style=textColor)

            if DRAW_FEATURES:

                for i in self.features:
                    window.deleteObject(i)

                self.features = [window.drawCircle((p.x, p.y), 1, style=featureStyle) for p in fsdkTracker.GetFacialFeatures(0, self.face_id)]

            return True

        except FSDK.IdNotFound:
            pass

        return False


trackers = {}
mouse_pos = None
onLivenessButton = None
liveness_status = False


def onMouseMove(event):
    global mouse_pos, onLivenessButton
    mouse_pos = event.x, event.y
    onLivenessButton = event.x >= livenessButton[0] and event.y >= livenessButton[1] and event.x < livenessButton[2] and event.y < livenessButton[3]

def onMousePress(event):
    if mouse_pos is not None:
        if onLivenessButton:
            global liveness_status
            liveness_status = not liveness_status
            fsdkTracker.SetParameter('DetectLiveness', liveness_status)
            return
        for face in trackers.values():
            if face.frame is not None and face.frame.contains(*mouse_pos):
                name = tkinter.simpledialog.askstring('Edit name', 'Enter person\'s name', initialvalue=face.name)
                if name:
                    fsdkTracker.LockID(face.face_id)
                    fsdkTracker.SetName(face.face_id, name)
                    fsdkTracker.UnlockID(face.face_id)
                    face.name = name
                return

window.canvas.bind('<Button-1>', onMousePress)
window.canvas.bind('<Motion>', onMouseMove)

try:
    while not need_to_exit:
        img = camera.GrabFrame()
        img = img.Resize(window.getScaleFor(img.width, img.height))
        window.drawImage(Image.frombytes('RGB', (img.width, img.height), img.ToBytes(FSDK.FSDK_IMAGE_COLOR_24BIT)))
        
        window.drawRectangle(*livenessButton, style = livenessActiveStyle if onLivenessButton else livenessStyle)
        liveveness = f"Liveness detection: {'On' if liveness_status else 'Off'}"
        window.drawText((livenessButton[2]//2, livenessButton[3]//2), liveveness, style=textColor)

        faces = frozenset(fsdkTracker.FeedFrame(0, img)) # recognize all faces in the image
        for face_id in faces.difference(trackers.keys()): # create new trackers
            trackers[face_id] = FaceLocator(face_id)

        img.Free()

        missed = []
        for face_id, tracker in trackers.items(): # iterate over current trackers
            if not face_id in faces:
                missed.append(face_id)
                continue

            tracker.draw()

        for face_id in missed:
            del trackers[face_id]
        window.update_idletasks()
        window.update()
finally:
    print("Please wait while freeing resources... ",  flush=True)
    fsdkTracker.SaveToFile(trackerMemoryFile)

root.destroy()

fsdkTracker.Free()
camera.Close()

FSDK.FinalizeCapturing()

FSDK.Finalize()
