# USAGE
# python detect_mask_video.py

# import the necessary packages
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.models import load_model
import numpy as np
import argparse
import os
import cv2
import sys
from PyQt5 import QtCore
from PyQt5 import QtWidgets
from PyQt5 import QtGui
from PyQt5 import QtTest
import pyaudio
import wave
import requests
import time

#Audio
# Record Audio의 startRecording 메서드에서 input_device_index는 기기마다 다름.
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16000
CHUNK = 1024
MAX_RECORD_SECONDS = 30
WAVE_OUTPUT_FILENAME = "saved_voice\\audiofile\\file.wav"
WAVE_ENROLL_FILENAME = "saved_voice\\enrollfile\\file.wav"
#URL
URL = 'http://163.180.146.68:7777/{}'
#SpeakerRecognition
THRESHOLD = 0.8
SPEAKER_ID = 'NA'

class ShowVideo(QtCore.QObject):
    flag_detect_mask = True
    run_video = True

    camera = cv2.VideoCapture(0)  # 연결된 영상장치 index, 기본은 0

    ret, image = camera.read()  # 2개의 값 리턴, 첫 번째는 프레임 읽음여부, 두 번째는 프레임 자체
    height, width = image.shape[:2]

    VideoSignal1 = QtCore.pyqtSignal(QtGui.QImage)  # VideoSignal1이라는 사용자 정의 시그널 생성

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

    @QtCore.pyqtSlot()
    def startVideo(self, faceNet, maskNet):
        global image

        run_video = True
        self.flag_detect_mask = True
        while run_video:
            ret, image = self.camera.read()

            # detect faces in the frame and determine if they are wearing a
            # face mask or not
            QtWidgets.QApplication.processEvents()
            if self.flag_detect_mask:
                (locs, preds) = detect_and_predict_mask(image, faceNet, maskNet)

                # QtWidgets.QApplication.processEvents()
            # if self.flag_detect_mask:
                frame = image
                # loop over the detected face locations and their corresponding
                # locations
                for (box, pred) in zip(locs, preds):
                    # unpack the bounding box and predictions
                    (startX, startY, endX, endY) = box
                    (mask, withoutMask) = pred

                    # determine the class label and color we'll use to draw
                    # the bounding box and text
                    label = "Mask" if mask > withoutMask else "No Mask"  # 박스 상단 출력 string
                    color = (0, 255, 0) if label == "Mask" else (0, 0, 255)

                    # include the probability in the label
                    label = "{}: {:.2f}%".format(label, max(mask, withoutMask) * 100)

                    # display the label and bounding box rectangle on the output
                    # frame
                    cv2.putText(frame, label, (startX, startY - 10),  # label에 string들어감
                                cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 2)
                    cv2.rectangle(frame, (startX, startY), (endX, endY), color, 2)
                image = frame
            ###
            color_swapped_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

            qt_image1 = QtGui.QImage(color_swapped_image.data,
                                     self.width,
                                     self.height,
                                     color_swapped_image.strides[0],
                                     QtGui.QImage.Format_RGB888)
            self.VideoSignal1.emit(qt_image1)

            loop = QtCore.QEventLoop()
            QtCore.QTimer.singleShot(25, loop.quit)  # 25 ms
            loop.exec_()

    @QtCore.pyqtSlot()
    def maskdetectionoff(self):
        self.flag_detect_mask = False


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

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

    def initUI(self):
        self.setWindowTitle('Webcam')

    @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()


def detect_and_predict_mask(frame, faceNet, maskNet):
    # grab the dimensions of the frame and then construct a blob
    # from it
    (h, w) = frame.shape[:2]
    blob = cv2.dnn.blobFromImage(frame, 1.0, (300, 300),
                                 (104.0, 177.0, 123.0))

    # pass the blob through the network and obtain the face detections
    faceNet.setInput(blob)
    detections = faceNet.forward()

    # initialize our list of faces, their corresponding locations,
    # and the list of predictions from our face mask network
    faces = []
    locs = []
    preds = []

    # loop over the detections
    for i in range(0, detections.shape[2]):
        # extract the confidence (i.e., probability) associated with
        # the detection
        confidence = detections[0, 0, i, 2]

        # filter out weak detections by ensuring the confidence is
        # greater than the minimum confidence
        if confidence > args["confidence"]:
            # compute the (x, y)-coordinates of the bounding box for
            # the object
            box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
            (startX, startY, endX, endY) = box.astype("int")

            # ensure the bounding boxes fall within the dimensions of
            # the frame
            (startX, startY) = (max(0, startX), max(0, startY))
            (endX, endY) = (min(w - 1, endX), min(h - 1, endY))

            # extract the face ROI, convert it from BGR to RGB channel
            # ordering, resize it to 224x224, and preprocess it
            face = frame[startY:endY, startX:endX]
            face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
            face = cv2.resize(face, (224, 224))
            face = img_to_array(face)
            face = preprocess_input(face)

            # add the face and bounding boxes to their respective
            # lists
            faces.append(face)
            locs.append((startX, startY, endX, endY))

    # only make a predictions if at least one face was detected
    if len(faces) > 0:
        # for faster inference we'll make batch predictions on *all*
        # faces at the same time rather than one-by-one predictions
        # in the above `for` loop
        faces = np.array(faces, dtype="float32")
        preds = maskNet.predict(faces, batch_size=32)

    # return a 2-tuple of the face locations and their corresponding
    # locations
    return (locs, preds)


class SpeakerRecognition(QtWidgets.QWidget):
    verification_url = URL.format('verification')
    identification_url = URL.format('identification')
    enrollment_url = URL.format('enroll')
    speaker_id = ''

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

    def initUI(self):
        self.label_1_1 = QtWidgets.QLabel('Result Message: ', self)
        self.label_1_2 = QtWidgets.QLabel('', self)
        self.push_button5 =QtWidgets.QPushButton('Authenticate', self)
        self.push_button5.clicked.connect(self.doAction)

        self.dialog_button = QtWidgets.QPushButton('화자 ID 입력:', self)
        self.dialog_button.clicked.connect(self.showDialog)
        self.le = QtWidgets.QLineEdit(self)

        self.register_button = QtWidgets.QPushButton('Register new voice', self)
        self.register_button.clicked.connect(self.switch_enrollment)

    def verification(self, speaker):
        try:
            with open(WAVE_OUTPUT_FILENAME, 'rb') as file_opened:
                files = {'file': file_opened}
                data = {'enroll_speaker': speaker}
                r = requests.post(self.verification_url, files=files, data=data)
                print(r.text)
                return r.text
        except FileNotFoundError:
            return False

    def identification(self):
        try:
            with open(WAVE_OUTPUT_FILENAME, 'rb') as file_opened:
                files = {'file': file_opened}
                r = requests.post(self.identification_url, files=files)
                print(r.text)
                return r.text
        except FileNotFoundError:
            return False

    def recognition(self):
        speaker = self.identification()
        if speaker == False:
            print('Record voice first!')
            return False

        percentage = self.verification(speaker)
        print(speaker, percentage)

        if float(percentage) >= THRESHOLD:
            result = '승인! 등록된 화자입니다.(ID: {})'.format(speaker)
            #result = speaker
        else:
            result = '등록되지 않은 화자입니다!'
        return result

    @QtCore.pyqtSlot()
    def doAction(self):
        start = time.time()
        recog = self.recognition()
        print("Inference time : ", time.time() - start)
        if recog == False:
            self.label_1_2.setText('Voice not recorded, record voice first!')
        else:
            self.label_1_2.setText(recog)

    def enrollment(self, speaker_id):
        try:
            if speaker_id == '':
                return 0
            with open(WAVE_ENROLL_FILENAME, 'rb') as file_opened:
                files = {'file': file_opened}
                data = {'enroll_speaker': speaker_id}
                r = requests.post(self.enrollment_url, files=files, data=data)
                print(r.text)
                return r.text
        except FileNotFoundError:
            return 1

    def switch_enrollment(self):
        enroll = self.enrollment(self.speaker_id)
        if enroll == 1:
            self.label_1_2.setText('Voice not recorded, record voice first!')
        elif enroll == 0:
            self.label_1_2.setText('No speaker ID input!')
        else:
            self.label_1_2.setText("""New speaker registered!('%s')""" % self.speaker_id)
            self.speaker_id = ''
            self.le.setText(self.speaker_id)

    def showDialog(self):
        text, ok = QtWidgets.QInputDialog.getText(self, '화자 등록',
                                                  '등록할 화자 ID(Unique 값)을 입력하십시오:')
        if ok:
            self.le.setText(str(text))
            self.speaker_id = str(text)

class RecordAudio(QtCore.QObject):
    isrecording = False
    frames = []

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

    @QtCore.pyqtSlot()
    def startRecording(self):
        # start Recording
        self.audio = pyaudio.PyAudio()
        self.stream = self.audio.open(format=pyaudio.paInt16,
                                      channels=CHANNELS,
                                      rate=RATE,
                                      input=True,
                                      input_device_index=1,  # 기기마다 마이크 인덱스 다름
                                      frames_per_buffer=CHUNK)
        self.isrecording = True
        print("recording...")

        # frames = []
        self.frames.clear()

        for i in range(0, int(RATE / CHUNK * MAX_RECORD_SECONDS)):
            QtWidgets.QApplication.processEvents()
            if self.isrecording:
                data = self.stream.read(CHUNK)
                self.frames.append(data)
            else:
                print("Stopped recording")
                break
        print("finished recording")

        # stop Recording
        self.stream.stop_stream()
        self.stream.close()
        self.audio.terminate()
        waveFile = wave.open(WAVE_OUTPUT_FILENAME, 'wb')
        waveFile.setnchannels(CHANNELS)
        waveFile.setsampwidth(self.audio.get_sample_size(FORMAT))
        waveFile.setframerate(RATE)
        waveFile.writeframes(b''.join(self.frames))
        waveFile.close()
        self.frames.clear()

    def stopRecording(self):
        print("stop called")
        self.isrecording = False

    def switch(self):
        if self.isrecording:
            QtTest.QTest.qWait(1 * 1000)
            self.stopRecording()
        else:
            self.startRecording()

class RecordAudio_enroll(QtCore.QObject):
    isrecording = False
    frames = []

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

    @QtCore.pyqtSlot()
    def startRecording(self):
        # start Recording
        self.audio = pyaudio.PyAudio()
        self.stream = self.audio.open(format=pyaudio.paInt16,
                                      channels=CHANNELS,
                                      rate=RATE,
                                      input=True,
                                      input_device_index=1,  # 기기마다 마이크 인덱스 다름
                                      frames_per_buffer=CHUNK)
        self.isrecording = True
        print("recording...")

        # frames = []
        self.frames.clear()

        for i in range(0, int(RATE / CHUNK * MAX_RECORD_SECONDS)):
            QtWidgets.QApplication.processEvents()
            if self.isrecording:
                data = self.stream.read(CHUNK)
                self.frames.append(data)
            else:
                print("Stopped recording")
                break
        print("finished recording")

        # stop Recording
        self.stream.stop_stream()
        self.stream.close()
        self.audio.terminate()
        waveFile = wave.open(WAVE_ENROLL_FILENAME, 'wb')
        waveFile.setnchannels(CHANNELS)
        waveFile.setsampwidth(self.audio.get_sample_size(FORMAT))
        waveFile.setframerate(RATE)
        waveFile.writeframes(b''.join(self.frames))
        waveFile.close()
        self.frames.clear()

    def stopRecording(self):
        print("stop called")
        self.isrecording = False

    def switch(self):
        if self.isrecording:
            QtTest.QTest.qWait(1 * 1000)
            self.stopRecording()
        else:
            self.startRecording()

class RecordViewer(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super(RecordViewer, self).__init__(parent)
        self.initUI()

    def initUI(self):
        self.pbar = QtWidgets.QProgressBar(self)
        self.pbar.setFixedWidth(400)
        self.pbar.setMaximum(MAX_RECORD_SECONDS)
        self.pbar.setAlignment(QtCore.Qt.AlignCenter)

        self.push_button3 = QtWidgets.QPushButton('Start Audio Record', self)
        self.push_button3.clicked.connect(self.doAction)

        self.timer = QtCore.QBasicTimer()
        self.step = 0

    def timerEvent(self, e):
        if self.step >= MAX_RECORD_SECONDS:
            self.timer.stop()
            self.push_button3.setText("Restart")
            return
        self.step = self.step + 1
        self.pbar.setValue(self.step)
        self.pbar.setFormat("%d sec" % self.step)

    @QtCore.pyqtSlot()
    def doAction(self):
        if self.timer.isActive():
            self.timer.stop()
            self.push_button3.setText("Restart")
        else:
            self.pbar.reset()
            self.step = 0
            self.timer.start(1000, self)  # 1000/1000초마다 timer실행
            self.push_button3.setText("Stop")


if __name__ == '__main__':
    # construct the argument parser and parse the arguments
    ap = argparse.ArgumentParser()
    ap.add_argument("-f", "--face", type=str, default="face_detector",
                    help="path to face detector model directory")
    ap.add_argument("-m", "--model", type=str, default="mask_detector.model",
                    help="path to trained face mask detector model")
    ap.add_argument("-c", "--confidence", type=float, default=0.5,
                    help="minimum probability to filter weak detections")
    args = vars(ap.parse_args())

    # load our serialized face detector model from disk
    print("[INFO] loading face detector model...")
    prototxtPath = os.path.sep.join([args["face"], "deploy.prototxt"])
    weightsPath = os.path.sep.join([args["face"],
                                    "res10_300x300_ssd_iter_140000.caffemodel"])
    faceNet = cv2.dnn.readNet(prototxtPath, weightsPath)

    # load the face mask detector model from disk
    print("[INFO] loading face mask detector model...")
    maskNet = load_model(args["model"])

    app = QtWidgets.QApplication(sys.argv)  # app 생성

    thread = QtCore.QThread()
    thread.start()
    vid = ShowVideo()
    vid.moveToThread(thread)

    thread2 = QtCore.QThread()
    thread2.start()
    aud = RecordViewer()
    aud.moveToThread(thread2)

    thread3 = QtCore.QThread()
    thread3.start()
    mic = RecordAudio_enroll()
    mic.moveToThread(thread3)

    thread4 = QtCore.QThread()
    thread4.start()
    sr = SpeakerRecognition()
    sr.moveToThread(thread4)

    thread5 = QtCore.QThread()
    thread5.start()
    aud2 = RecordViewer()
    aud2.moveToThread(thread5)

    thread6 = QtCore.QThread()
    thread6.start()
    mic2 = RecordAudio()
    mic2.moveToThread(thread6)

    thread7 = QtCore.QThread()
    thread7.start()
    sr2 = SpeakerRecognition()
    sr2.moveToThread(thread7)


    image_viewer1 = ImageViewer()

    vid.VideoSignal1.connect(image_viewer1.setImage)

    push_button1 = QtWidgets.QPushButton('Start Mask Detection')
    push_button2 = QtWidgets.QPushButton('Mask Detection Off')
    push_button4 = QtWidgets.QPushButton('Close')

    push_button1.clicked.connect(lambda: vid.startVideo(faceNet, maskNet))
    push_button2.clicked.connect(vid.maskdetectionoff)
    aud.push_button3.clicked.connect(mic.switch)
    push_button4.clicked.connect(sys.exit)
    aud2.push_button3.clicked.connect(mic2.switch)

    empty_label = QtWidgets.QLabel()
    empty_label.setText('')

    L_groupBox = QtWidgets.QGroupBox("Mask Detection")
    LR_layout = QtWidgets.QVBoxLayout()
    LR_layout.addWidget(push_button1)
    LR_layout.addWidget(push_button2)
    LR_layout.addStretch(1)

    L_horizontal_layout1 = QtWidgets.QHBoxLayout()
    L_horizontal_layout1.addWidget(image_viewer1)
    L_horizontal_layout1.addLayout(LR_layout)
    L_groupBox.setLayout(L_horizontal_layout1)

    RU_groupBox = QtWidgets.QGroupBox("Speaker Registration")
    pbar_layout = QtWidgets.QHBoxLayout()
    pbar_layout.addWidget(aud.pbar)
    pbar_layout.addStretch(1)
    ##
    dialog_layout = QtWidgets.QHBoxLayout()
    dialog_layout.addWidget(sr2.dialog_button)
    dialog_layout.addWidget(sr2.le)
    dialog_layout.addStretch(1)

    register_layout = QtWidgets.QHBoxLayout()
    register_layout.addWidget(sr2.register_button)

    result_1_layout = QtWidgets.QHBoxLayout()
    result_1_layout.addWidget(sr2.label_1_1)
    result_1_layout.addWidget(sr2.label_1_2)
    result_1_layout.addStretch(1)
    ##
    RL_label1 = QtWidgets.QLabel()
    RL_label1.setText("Max Record Time: 30 sec")
    RL_label2 = QtWidgets.QLabel()
    RL_label2.setText("Press Start/Restart to begin recording")

    RL_layout = QtWidgets.QVBoxLayout()
    RL_layout.addLayout(pbar_layout)
    RL_layout.addWidget(RL_label1)
    RL_layout.addWidget(RL_label2)
    RL_layout.addLayout(dialog_layout)
    RL_layout.addLayout(result_1_layout)
    RL_layout.addStretch(1)

    push_button3_layout = QtWidgets.QHBoxLayout()
    push_button3_layout.addWidget(aud.push_button3)
    # push_button3_layout.addStretch(1)

    # close_layout = QtWidgets.QHBoxLayout()
    # close_layout.addWidget(push_button4)

    RR_layout = QtWidgets.QVBoxLayout()
    RR_layout.addLayout(push_button3_layout)
    RR_layout.addWidget(empty_label)
    RR_layout.addWidget(empty_label)
    RR_layout.addLayout(register_layout)
    RR_layout.addStretch(1)
    # RR_layout.addLayout(close_layout)

    R_horizontal_layout2 = QtWidgets.QHBoxLayout()
    R_horizontal_layout2.addLayout(RL_layout)
    R_horizontal_layout2.addLayout(RR_layout)
    RU_groupBox.setLayout(R_horizontal_layout2)


    RD_groupBox = QtWidgets.QGroupBox("Speaker Recognition")
###
    pbar_2_layout = QtWidgets.QHBoxLayout()
    pbar_2_layout.addWidget(aud2.pbar)
    pbar_2_layout.addStretch(1)

    RDL_label1 = QtWidgets.QLabel()
    RDL_label1.setText("Max Record Time: 30 sec")
    RDL_label2 = QtWidgets.QLabel()
    RDL_label2.setText("Press Start/Restart to begin recording")

    push_button3_2_layout = QtWidgets.QHBoxLayout()
    push_button3_2_layout.addWidget(aud2.push_button3)
###
    result_2_layout = QtWidgets.QHBoxLayout()
    result_2_layout.addWidget(sr.label_1_1)
    result_2_layout.addWidget(sr.label_1_2)
    result_2_layout.addStretch(1)

    RDL_layout = QtWidgets.QVBoxLayout()
    RDL_layout.addLayout(pbar_2_layout)
    RDL_layout.addWidget(RDL_label1)
    RDL_layout.addWidget(RDL_label2)
    RDL_layout.addWidget(empty_label)
    RDL_layout.addLayout(result_2_layout)
    RDL_layout.addStretch(1)

    push_button5_layout = QtWidgets.QHBoxLayout()
    push_button5_layout.addWidget(sr.push_button5)

    close_layout = QtWidgets.QHBoxLayout()
    close_layout.addWidget(push_button4)

    RDR_layout = QtWidgets.QVBoxLayout()
    RDR_layout.addLayout(push_button3_2_layout)
    RDR_layout.addWidget(empty_label)
    RDR_layout.addWidget(empty_label)
    RDR_layout.addWidget(empty_label)
    RDR_layout.addLayout(push_button5_layout)
    RDR_layout.addStretch(1)
    RDR_layout.addLayout(close_layout)

    RD_horizontal_layout = QtWidgets.QHBoxLayout()
    RD_horizontal_layout.addLayout(RDL_layout)
    RD_horizontal_layout.addLayout(RDR_layout)
    RD_groupBox.setLayout(RD_horizontal_layout)

    R_layout = QtWidgets.QVBoxLayout()
    R_layout.addWidget(RU_groupBox)
    R_layout.addWidget(RD_groupBox)

    layout = QtWidgets.QHBoxLayout()
    layout.addWidget(L_groupBox)
    layout.addLayout(R_layout)

    layout_widget = QtWidgets.QWidget()
    layout_widget.setLayout(layout)

    main_window = QtWidgets.QMainWindow()
    main_window.setGeometry(150, 150, 500, 500)  # test
    main_window.setCentralWidget(layout_widget)
    main_window.setWindowTitle('마스크 디텍션 및 화자 식별을 통한 입출입 시스템')  # main window 제목
    main_window.show()
    sys.exit(app.exec_())  # 프로그램 대기상태 유지, 무한루프