# 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

FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16000
CHUNK = 1024
MAX_RECORD_SECONDS = 30
WAVE_OUTPUT_FILENAME = "audiofile\\file.wav"
#Record Audio의 startRecording 메서드에서 input_device_index는 기기마다 다름.

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

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

    # test
    thread3 = QtCore.QThread()
    thread3.start()
    mic = RecordAudio()
    mic.moveToThread(thread3)

    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)

    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)

    R_groupBox = QtWidgets.QGroupBox("Audio Record")
    pbar_layout = QtWidgets.QHBoxLayout()
    pbar_layout.addWidget(aud.pbar)
    pbar_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.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.addStretch(1)
    RR_layout.addLayout(close_layout)

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

    layout = QtWidgets.QHBoxLayout()
    layout.addWidget(L_groupBox)
    layout.addWidget(R_groupBox)

    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_())   # 프로그램 대기상태 유지, 무한루프