################################################################################
#####################    Autor: Matúš Juhász  ##################################

# This program is free software: you can redistribute it and/or modify
#     it under the terms of the GNU General Public License as published by
#     the Free Software Foundation, either version 3 of the License, or
#     (at your option) any later version.
#
#     This program is distributed in the hope that it will be useful,
#     but WITHOUT ANY WARRANTY; without even the implied warranty of
#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#     GNU General Public License for more details.
#
#     You should have received a copy of the GNU General Public License
#     along with this program.  If not, see <https://www.gnu.org/licenses/>.

import cv2
import numpy as np
import sys
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
import time

import matplotlib.pyplot as plt
import threading

import math
import os
import dlib

from PyQt5 import QtCore, QtGui, QtWidgets, QtMultimedia, QtMultimediaWidgets
from PyQt5.QtCore import QSize, pyqtSignal, Qt
from PyQt5.QtMultimedia import QMediaPlayer
from PyQt5.QtMultimediaWidgets import QVideoWidget
from PyQt5.QtWidgets import QApplication, QWidget, QInputDialog, QLineEdit, QFileDialog, QSlider, QHBoxLayout, QVBoxLayout, QMessageBox, QCheckBox, QGroupBox

from matplotlib.colors import Normalize
from os import listdir
from os.path import isfile, join
from pathlib import Path

from TF_face_detector.utils import label_map_util
from TF_face_detector.utils import visualization_utils_color as vis_util

# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_CKPT = './TF_face_detector/frozen_inference_graph_face.pb'

# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = './TF_face_detector/face_label_map.pbtxt'

NUM_CLASSES = 2

label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)

class Face_Detector_Haar:
    def check_image_haar(self, img):
        face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        face = face_cascade.detectMultiScale(gray,
                                            scaleFactor=1.1,
                                            minNeighbors=2
                                            )

        for (x,y,w,h) in face:
            #vykreslenie obldznika okolo hornej casti tela
            cv2.rectangle(img, (x,y), (x+w,y+h), (0,255,0), 2)

        return face
# kod triedy ciastocne prevzaty z https://github.com/yeephycho/tensorflow-face-detection
class DetectorFace(object):
    def __init__(self, PATH_TO_CKPT):
        # self.path_to_ckpt = path_to_ckpt

        self.detection_graph = tf.Graph()
        with self.detection_graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')

        with self.detection_graph.as_default():
            config = tf.ConfigProto()
            config.gpu_options.allow_growth = True
            self.sess = tf.Session(graph=self.detection_graph, config=config)

        self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
        # Each box represents a part of the image where a particular object was detected.
        self.boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
        # Each score represent how level of confidence for each of the objects.
        # Score is shown on the result image, together with the class label.
        self.scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
        self.classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
        self.num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')

    def load_image_into_numpy_array(self, image):
        (im_width, im_height) = image.size
        return np.array(image.getdata()).reshape(
            (im_height, im_width, 3)).astype(np.uint8)

    def process(self, image):
        image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # the array based representation of the image will be used later in order to prepare the
        # result image with boxes and labels on it.
        # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
        image_np_expanded = np.expand_dims(image_np, axis=0)
        # Actual detection.
        start_time = time.time()

        (boxes, scores, classes, num_detections) = self.sess.run(
            [self.boxes, self.scores, self.classes, self.num_detections],
            feed_dict={self.image_tensor: image_np_expanded})

        elapsed_time = time.time() - start_time
        print('inference time cost: {}'.format(elapsed_time))

        return (boxes, scores, classes, num_detections)

    def close(self):
        self.sess.close()

# kod triedy ciastocne prevzaty z https://gist.github.com/madhawav/1546a4b99c8313f06c0b2d7d7b4a09e2
class DetectorAPI:
    def __init__(self, path_to_ckpt):
        self.path_to_ckpt = path_to_ckpt

        self.detection_graph = tf.Graph()
        with self.detection_graph.as_default():
            od_graph_def = tf.compat.v1.GraphDef()
            with tf.compat.v2.io.gfile.GFile(self.path_to_ckpt, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')

        self.default_graph = self.detection_graph.as_default()
        self.sess = tf.Session(graph=self.detection_graph)

        # Definite input and output Tensors for detection_graph
        self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
        # Each box represents a part of the image where a particular object was detected.
        self.detection_boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
        # Each score represent how level of confidence for each of the objects.
        # Score is shown on the result image, together with the class label.
        self.detection_scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
        self.detection_classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
        self.num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')

    def processFrame(self, image):
        # Expand dimensions since the trained_model expects images to have shape: [1, None, None, 3]
        image_np_expanded = np.expand_dims(image, axis=0)
        # Actual detection.
        (boxes, scores, classes, num) = self.sess.run(
            [self.detection_boxes, self.detection_scores, self.detection_classes, self.num_detections],
            feed_dict={self.image_tensor: image_np_expanded})

        # print("Elapsed Time:", end_time-start_time)
        im_height, im_width,_ = image.shape
        boxes_list = [None for i in range(boxes.shape[1])]
        for i in range(boxes.shape[1]):
            boxes_list[i] = (int(boxes[0,i,0] * im_height),
                        int(boxes[0,i,1]*im_width),
                        int(boxes[0,i,2] * im_height),
                        int(boxes[0,i,3]*im_width))

        return boxes_list, scores[0].tolist(), [int(x) for x in classes[0].tolist()], int(num[0])

    def close(self):
        self.sess.close()
        # self.default_graph.

class Heatmapclass(Normalize):
    time_list = []
    frame_list = []
    time_list_ssd = []
    frame_list_ssd = []
    # 72 reprezentuje 720x1280 cize stlpce x riadky
    arr = np.zeros((72,128), dtype=int)

    def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):
        self.midpoint = midpoint
        Normalize.__init__(self, vmin, vmax, clip)

    def __call__(self, value, clip=None):
        # I'm ignoring masked values and all kinds of edge cases to make a
        x, y = [self.vmin, self.midpoint, self.vmax], [0, 0.5, 1]
        return np.ma.masked_array(np.interp(value, x, y))

    def collect(self, x, y, flag):
        if flag:
            Heatmapclass.arr[y, x] += 1
        else:
            Heatmapclass.arr[y, x] -= 1

    def clean(self):
        Heatmapclass.arr = np.zeros((72,128), dtype=int)
        Heatmapclass.time_list = []
        Heatmapclass.frame_list = []

    def show_heatmap(self):
        threading.Thread(target=self.heat_map).start()
        
    def heat_map(self):
        w, h = plt.figaspect(Heatmapclass.arr)
        fig, ax = plt.subplots(figsize=(w, h))
        fig.subplots_adjust(left=0.07, bottom=0.05, right=1, top=0.95)
        norm = Heatmapclass(midpoint=0)
        im = ax.imshow(Heatmapclass.arr, norm=norm, cmap=plt.cm.seismic, interpolation='none', extent=[0, 128, 0, 72])
        fig.colorbar(im)
        figure_string = "Time when Heatmap captured: " + self.time()
        fig.canvas.setWindowTitle(figure_string)

        plt.title('Heatmap example')
        plt.ylabel('Y')
        plt.xlabel('X')
        plt.show()
        
        #close heatmap inmediatly
        #plt.close()

    def time(self):
        video = ShowVideo()
        if video.videotime is None:
            video_time = "0.00 sec"
        elif (video.videotime / 60) < 1:
            video_time = str(video.videotime) + " sec"
        else:
            video_time = str(video.videotime) + " min"

        return video_time

    def graf_collect(self, time, frame, bool):
        if bool:
            Heatmapclass.time_list_ssd.append(time)
            Heatmapclass.frame_list_ssd.append(frame)
        else:
            Heatmapclass.time_list.append(time)
            Heatmapclass.frame_list.append(frame)


    def graf_plot(self):
        fig, ax1 = plt.subplots()
        plt.xlabel('Frame')
        plt.ylabel('Time')
        plt.title('Speed Graph')
        fig.subplots_adjust(left=0.05, bottom=0.05, right=0.95, top=0.95)
        plt.grid(True)
        plt.xticks(np.arange(min(Heatmapclass.frame_list[1:]), max(Heatmapclass.frame_list[1:]), 5))
        plt.yticks(np.arange(min(Heatmapclass.time_list[1:]), max(Heatmapclass.time_list[1:]), 0.1))
        plt.xlim(min(Heatmapclass.frame_list[1:]))

        ax2 = ax1.twinx()
        ax1.plot(Heatmapclass.frame_list[1:], Heatmapclass.time_list[1:], 'g-')
        ax2.plot(Heatmapclass.frame_list_ssd[1:], Heatmapclass.time_list_ssd[1:], 'r-')
        plt.tight_layout()
        plt.show()

class ShowVideo(QtCore.QObject):
    videotime = None
    loop = None

    def __init__(self, parent = None):
        super(ShowVideo, self).__init__(parent)

    #initiating the built in camera
    camera_port = 0
    VideoSignal = QtCore.pyqtSignal(QtGui.QImage)

    @QtCore.pyqtSlot()
    def startVideo(self):

        # create class
        video = VideoWindow()
        heat = Heatmapclass()

        heat.clean()


        # set path to video file
        path = video.filename

        # Path for now TODO change
        model_path = 'frozen_inference_graph.pb'
        odapi = DetectorAPI(path_to_ckpt=model_path)
        threshold = 0.7

        odapiF = DetectorFace(PATH_TO_CKPT)
        # Class for detection face from body using HaarCasade algorithm
        detect_face_haar = Face_Detector_Haar()

        # set video input
        camera = cv2.VideoCapture(path)

        # frame count
        number_frame = 0
        fps = camera.get(cv2.CAP_PROP_FPS)

        # while is video opened
        loop = camera.isOpened()
        while loop:
            if video.end_while:
                break
            # read every frame
            ret, image = camera.read()
            if ret != True:
                print("CHYBA")
                break;
            # resize for faster computing
            image = cv2.resize(image, (1280, 720))
            # start time
            start_time = time.time()

            # Debug purposes numbering frames
            number_frame += 1
            print("frame number: ", number_frame)

            ShowVideo.videotime = number_frame / fps

            # ulozi suradnice body obdlznika do boxes
            boxes, scores, classes, num = odapi.processFrame(image)

            # procces face for every frame:
            (boxes_SSD, scores_SSD, classes_SSD, num_detections_SSD) = odapiF.process(image)
            boxes_SSD = np.squeeze(boxes_SSD)
            classes_SSD = np.squeeze(classes_SSD).astype(np.int32)
            scores_SSD = np.squeeze(scores_SSD)
            face_list = [None for pokus in range(min(20, boxes_SSD.shape[0]))]
            for pokus in range(min(20, boxes_SSD.shape[0])):
                if scores_SSD is None or scores_SSD[pokus] > .7:
                    box_face = tuple(boxes_SSD[pokus].tolist())
                    ymin, xmin, ymax, xmax = box_face
                    im_height, im_width,_ = image.shape
                    # create listdi
                    face_list[pokus] = (int(xmin * im_width), int(xmax * im_width),
                                                 int(ymin * im_height), int(ymax * im_height))

            body = 0
            for i in range(len(boxes)):
                if classes[i] == 1 and scores[i] > threshold:
                    body +=1
                    # print("TELO: ", body)
                    box = boxes[i]
                    x = box[1]
                    y = box[0]
                    xw = box[3]
                    yh = box[2]
                    w = xw - x
                    h = yh - y
                    cv2.rectangle(image, (x, y), (xw, yh), (255,0,0), 2)

                    #stred body
                    x_fin = int(x + (w/2))
                    y_fin = int(y + (h*0.25))

                    # copy of rectangled img
                    x_cor = box[0]
                    x_w = box[2]
                    y_cor = box[1]
                    y_h = box[3]
                    if (x_cor > 0 and y_cor > 0):
                        x_cor = box[0] - 15
                        x_w = box[2] + 15
                        y_cor = box[1] - 15
                        y_h = box[3] + 15
                        if (x_cor < 0):
                            x_cor = 0
                        if (y_cor < 0):
                            y_cor = 0
                        if (x_w < 0):
                            x_w = 0
                        if (y_h < 0):
                            y_h = 0
                    if video.flag:
                        face = detect_face_haar.check_image_haar(image[x_cor:x_w, y_cor:y_h])
                        if len(face) != 0:
                            for (x_up, y_up, w_up, h_up) in face:
                                #vykreslenie obldznika okolo hornej casti tela
                                x_fin = int(x + w_up/2)
                                y_fin = int(y + h_up/2)
                            heat.collect(int(x_fin/10), int(y_fin/10), True)
                        else:
                            heat.collect(int(x_fin/10), int(y_fin/10), False)
                    else:
                        # kontrol if body rectangle overlap face rectangle, if yes return to heat map True else return False
                        face_is_in_body = False
                        face_list = list(filter(None, face_list))
                        for face_num in range(len(face_list)):
                            face_box = face_list[face_num]
                            left = face_box[0]
                            right = face_box[1]
                            top = face_box[2]
                            bottom = face_box[3]
                            if ((left < xw) and (right > x) and (top < yh) and (bottom > y)):
                                face_is_in_body = True
                                break;
                        # if overlap give coordinates to set into heatmap
                        if face_is_in_body:
                            w = right - left
                            h = bottom - top
                            x_fin = int(left + (w/2))
                            y_fin = int(top + (h/2))
                            heat.collect(int(x_fin/10), int(y_fin/10), True)
                        else:
                            # else aproxim coordinates of head
                            heat.collect(int(x_fin/10), int(y_fin/10), False)
                        # vizualization of head on video
                        vis_util.visualize_boxes_and_labels_on_image_array(
                            image,
                            boxes_SSD,
                            classes_SSD,
                            scores_SSD,
                            category_index,
                            use_normalized_coordinates=True,
                            line_thickness=4)

            color_swapped_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

            height, width, _ = color_swapped_image.shape

            qt_image = QtGui.QImage(color_swapped_image.data, width, height, color_swapped_image.strides[0], QtGui.QImage.Format_RGB888)

            self.VideoSignal.emit(qt_image)

        # delete class
        odapi.close()
        odapiF.close()

#       heat.show_heatmap()
#       print heatmap
#       heat.heat_map()
#       plt.close('all')

class ImageViewer(QtWidgets.QWidget):
    def __init__(self, parent = None):
        super(ImageViewer, self).__init__(parent)
        self.title = 'PyQt5 simple window'
        self.image = QtGui.QImage()
        self.setAttribute(QtCore.Qt.WA_OpaquePaintEvent)

        self.initUI()

    def paintEvent(self, event):
        painter = QtGui.QPainter(self)
        painter.drawImage(0,0, self.image)
        self.image = QtGui.QImage()

    def initUI(self):
        self.setWindowTitle(self.title)

    @QtCore.pyqtSlot(QtGui.QImage)
    def setImage(self, image):
        if image.isNull():
            print("Viewer Dropped frame!")

        self.image = image
        if image.size() != self.size():
            self.setFixedSize(image.size())
        self.update()

class VideoWindow(QtWidgets.QMainWindow):
    filename = ''
    flag = True
    end_while = None

    def __init__(self, parent=None):
        super(VideoWindow, self).__init__(parent)
        self.setWindowTitle("Video Player")
        self.resize(1300, 900)

        # buttons
        self.startButton = QtWidgets.QPushButton('Start')
        self.mapButton = QtWidgets.QPushButton('Map')
        self.stopButton = QtWidgets.QPushButton('Stop')
        # set to unclicable
        self.stopButton.setEnabled(False)
        self.startButton.setEnabled(False)
        self.startButton.clicked.connect(self.btnstate)

        # Checkbox area
        self.groupBox = QGroupBox("Choose Face recognition algorithm")
        self.groupBox.setFixedHeight(50)
        # Horizontal checkboxes
        self.horizontal_layout = QHBoxLayout()
        self.groupBox.setLayout(self.horizontal_layout)

        # name checkboxes
        self.haarCheckbox = QCheckBox("Haar", self)
        self.haarCheckbox.setChecked(True)
        self.hogCheckbox = QCheckBox("SSD mobilnet", self)
        self.horizontal_layout.addWidget(self.haarCheckbox)
        self.horizontal_layout.addWidget(self.hogCheckbox)
        # Change state of them
        self.haarCheckbox.stateChanged.connect(self.onStateChange)
        self.hogCheckbox.stateChanged.connect(self.onStateChange)

        # Create vertical layout
        self.vertical_layout = QtWidgets.QVBoxLayout()
        # Create horizntal layout for buttons
        self.horizontal_buttons = QHBoxLayout()
        # Add buttons to vertical layout
        self.horizontal_buttons.addWidget(self.startButton)
        self.horizontal_buttons.addWidget(self.stopButton)
        self.vertical_layout.addLayout(self.horizontal_buttons)
        self.vertical_layout.addWidget(self.mapButton)
        self.vertical_layout.addWidget(self.groupBox)

        self.layout_widget = QtWidgets.QWidget()
        self.layout_widget.setLayout(self.vertical_layout)

        self.setCentralWidget(self.layout_widget)

        # Create new action
        openAction = QtWidgets.QAction(QtGui.QIcon('open.png'), '&Open', self)
        openAction.setShortcut('Ctrl+O')
        openAction.triggered.connect(self.openFile)

        # Create exit action
        exitAction = QtWidgets.QAction(QtGui.QIcon('exit.png'), '&Exit', self)
        exitAction.setShortcut('Ctrl+Q')
        exitAction.triggered.connect(self.close)


        # Create menu bar and add action
        menuBar = self.menuBar()
        fileMenu = menuBar.addMenu('&Open File')
        fileMenu.addAction(openAction)
        fileMenu.addAction(exitAction)

        self.msgBox = QMessageBox()

    # change state of chcekbox if one is active
    @QtCore.pyqtSlot(int)
    def onStateChange(self, state):
        if state == Qt.Checked:
            if self.sender() == self.haarCheckbox:
                self.hogCheckbox.setChecked(False)
                VideoWindow.flag = True
            elif self.sender() == self.hogCheckbox:
                self.haarCheckbox.setChecked(False)
                VideoWindow.flag = False

    #open file function and enable play button
    def openFile(self):
        fileName, _ = QtWidgets.QFileDialog.getOpenFileName(self, "Open Movie")

        if fileName != '' :
            self.startButton.setEnabled(True)

        VideoWindow.filename = fileName

    # close function with dialog
    def closeEvent(self, event):
        self.msgBox.setText("Do you want close?")
        self.msgBox.setStandardButtons(QMessageBox.Cancel | QMessageBox.Yes | QMessageBox.No);
        self.msgBox.setDefaultButton(QMessageBox.Yes);
        if self.msgBox.exec_() == QMessageBox.Yes:
            event.accept()
        else:
            event.ignore()

    # enable play after stop
    def while_close(self, pressed):
        if not pressed:
            VideoWindow.end_while = True
            self.startButton.setEnabled(True)
            self.stopButton.setEnabled(False)

    def while_closesed(self):
        VideoWindow.end_while = True
        self.startButton.setEnabled(True)
        self.stopButton.setEnabled(False)

    # enable stop button after start
    def btnstate(self, pressed):
        if not pressed:
            VideoWindow.end_while = False
            self.startButton.setEnabled(False)
            self.stopButton.setEnabled(True)

if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    thread = QtCore.QThread()
    thread.start()
    # Video Class
    vid = ShowVideo()
    vid.moveToThread(thread)
    # Classes
    image_viewer = ImageViewer()
    map = Heatmapclass()
    # aplication = VideoWindow()

    vid.VideoSignal.connect(image_viewer.setImage)

    main_window = VideoWindow()

    def my_close():
        main_window.while_closesed()
        map.heat_map()

    # # Connect buttons with functions
    main_window.startButton.clicked.connect(vid.startVideo)
    main_window.mapButton.clicked.connect(map.heat_map)
    main_window.stopButton.clicked.connect(my_close)

    main_window.vertical_layout.addWidget(image_viewer)

    main_window.show()
    # Exit
    sys.exit(app.exec_())


