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

from PyQt5 import QtWidgets
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *

import NewUi_MainWindow as ui

from uvc_thermal import Thermal
from camera import Camera
from detector import FaceDetector
from anti import AntiSpoofing
from heater import Heater
from aligner import BoxAligner
from correcter import Correcter
from tracker import Tracker

import cv2
import numpy as np
import time

class Main(QMainWindow, ui.Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        #self.resize(640,480)
        self.showFullScreen()
        # 設置自動填滿大小
        self.frame.setScaledContents(True)
        self.thermal_frame.setScaledContents(True)
       
        self.normal_face1.setScaledContents(True)
        self.normal_face2.setScaledContents(True)
        self.normal_face3.setScaledContents(True)
        
        self.heigh_face1.setScaledContents(True)
        self.heigh_face2.setScaledContents(True)
        self.heigh_face3.setScaledContents(True)
        
        # 顯示人臉
        self.normal_faces = [self.normal_face1, self.normal_face2 ,self.normal_face3]
        self.heigh_faces = [self.heigh_face1, self.heigh_face2, self.heigh_face3]
        self.normal_face_count = 0
        self.heigh_face_count = 0
        
        # 載入相關元件
        self.thermal = Thermal(width=160, height=120, framerate=9, frame_width=160, frame_height=120, log=None)
        self.camera = Camera(width=640, height=480, framerate=30, log=None, channel=0, flip=None)
        self.detector = FaceDetector(model_path="./models/version-RFB/RFB-320.mnn" , input_size=(320,240))
        # self.anti = AntiSpoofing(model_path="../Ultra-Light-Fast-Generic-Face-Detector-1MB/MNN/model/4_0_0_80x80_MiniFASNetV1SE.mnn")
        self.aligner = BoxAligner(img_h = 480,img_w = 640)
        self.heater = Heater(pwm_pin=12 , target_temp=40.0)
        self.correcter = Correcter(model_path="../thermal-tk/20210421_correcter.sav")
        
        # Create Object Tracker
        self.tracker = Tracker(50, 10, 5, 0)
        # Variables initialization
        self.skip_frame_count = 0
        self.track_colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0),
                        (0, 255, 255), (255, 0, 255), (255, 127, 255),
                        (127, 0, 255), (127, 0, 127)]
        self.last_tracker_id = []
        
        self.thermal.start()
        time.sleep(2)
        self.camera.start()
        time.sleep(2)
        self.heater.start()
        time.sleep(2)
        self.num_frame = 0
        self.painter = QPainter(self)
#         self.multiTracker = cv2.MultiTracker_create()
#         self.success = False
    def paintEvent(self, QPaintEvent):
        # 繪圖事件
        ret, frame = self.camera.capture()
        thermal_ret,thermal_frame, thermal_row = self.thermal.capture()
        if(thermal_ret and ret):
            self.num_frame+=1
            blackbody_max = thermal_row[:,:34].max()
            heater_temp = self.heater.last_temp
            #heater_temp = 40.0
            # 偵測人臉
#             if(self.num_frame % 2 == 0):
            if(True):
                stime = time.time()
                boxes, labels, probs = self.detector.predict(frame)
                processing_time = time.time()-stime
                # 取出所有偵測的結果 並 加入追蹤
                centers = []
                shapes = []
                temperatures = []
                temperature_locs = []
                for i in range(boxes.shape[0]):
                    box = boxes[i, :]
                    score = probs[i]
                    if(score > 0.9):
                        thermal_box = self.aligner.box_aligment([box])[0]
                        x1 = max(box[0],0)
                        y1 = max(box[1],0)
                        x2 = min(box[2],frame.shape[1])
                        y2 = min(box[3],frame.shape[0])
                        centers.append([(x1+x2)/2,(y1+y2)/2])
                        shapes.append([(y2 - y1),(x2 - x1)]) # h,w
                        
                        face_temp = self.max_temperature(thermal_box,thermal_row,black_h=32)
                        corrected_face_temp = self.correcter.predict(x_detect_k= face_temp, x_black_k= blackbody_max, x_black_c= heater_temp)
                        #temperatures.append("{:.2f}->{:.2f}".format(face_temp*0.0092,corrected_face_temp))
                        temperatures.append(corrected_face_temp)
                        temperature_locs.append(thermal_box)
                if (len(centers) > 0):
                    # Track object using Kalman Filter
                    self.tracker.Update(centers, shapes, temperatures, temperature_locs)         
                    # For identified object tracks draw tracking line
                    # Use various colors to indicate different track_id
                    temp_tracker_id = []
                    for i in range(len(self.tracker.tracks)):
                        if (len(self.tracker.tracks[i].trace) > 1):
                            cx,cy = self.tracker.tracks[i].prediction[0]
                            tid = self.tracker.tracks[i].track_id
                            tx1 = cx - (self.tracker.tracks[0].w)/2
                            ty1 = cy - (self.tracker.tracks[0].h)/2
                            tx2 = cx + (self.tracker.tracks[0].w)/2
                            ty2 = cy + (self.tracker.tracks[0].h)/2
                            thermal_box = self.aligner.box_aligment([[tx1,ty1,tx2,ty2]])[0]
                            #cv2.rectangle(frame, (int(tx1), int(ty1)), (int(tx2), int(ty2)), track_colors[clr], 2)
                            color = (0, 255, 0)
                            face_frame = frame[y1:y2,x1:x2].copy()
                            face_frame=cv2.cvtColor(face_frame,cv2.COLOR_BGR2RGB)
                            Qface_frame=QImage(face_frame.data,face_frame.shape[1],face_frame.shape[0],face_frame.shape[1]*3,QImage.Format_RGB888)

                            face_pixMap=QPixmap.fromImage(Qface_frame)
                            
                            cv2.rectangle(thermal_frame, (thermal_box[0]//4, thermal_box[1]//4), (thermal_box[2]//4, thermal_box[3]//4), color, 2)
                            cv2.rectangle(frame, (int(tx1), int(ty1)), (int(tx2), int(ty2)), color, 2)
                            cv2.putText(frame, "{:.2f}".format(self.tracker.tracks[i].temp), (int(cx),int(cy)), cv2.FONT_HERSHEY_COMPLEX_SMALL,1, (0, 255, 255), 1, cv2.LINE_AA)
                            # 判斷溫度是否過高
                            temp_tracker_id.append(self.tracker.tracks[i].track_id)
                            if(not self.tracker.tracks[i].track_id in self.last_tracker_id):
                                if(self.tracker.tracks[i].temp < 37.0):
                                    self.normal_face_count+=1
                                    index = self.normal_face_count % 3
                                    self.normal_faces[index].setPixmap(face_pixMap)
                                else:
                                    self.heigh_face_count+=1
                                    index = self.heigh_face_count % 3
                                    self.heigh_faces[index].setPixmap(face_pixMap)
                                    
                    self.last_tracker_id = temp_tracker_id
                
                # 加熱片溫度顯示
                cv2.putText(thermal_frame, "{:.2f}".format(heater_temp), (320//4, 64//4), cv2.FONT_HERSHEY_COMPLEX_SMALL,1, (0, 0, 0), 1, cv2.LINE_AA)
            # 顯示至UI
            frame=cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
            thermal_frame=cv2.cvtColor(thermal_frame,cv2.COLOR_BGR2RGB)
            Qframe=QImage(frame.data,frame.shape[1],frame.shape[0],frame.shape[1]*3,QImage.Format_RGB888)
            Qthermal_frame=QImage(thermal_frame.data,thermal_frame.shape[1],thermal_frame.shape[0],thermal_frame.shape[1]*3,QImage.Format_RGB888)
            pixMap=QPixmap.fromImage(Qframe)
            thermal_pixMap=QPixmap.fromImage(Qthermal_frame)
            self.frame.setPixmap(pixMap)
            self.thermal_frame.setPixmap(thermal_pixMap)
            
            
            self.update()
    def closeEvent(self,event):
        # 關閉程式後執行
        self.thermal.stop()
        self.camera.stop()
        self.heater.stop()
    def max_temperature(self,box,thermal_row,black_h=0,thermal_height=120,thermal_width=160,image_height=480,image_width=640):
        scale_ratio_h = thermal_height/image_height
        scale_ratio_w = thermal_width/image_width
        x1 = max(int(box[0]*scale_ratio_w),0)
        y1 = max(int(box[1]*scale_ratio_h),black_h)
        x2 = min(int(box[2]*scale_ratio_w),thermal_width)
        y2 = min(int(box[3]*scale_ratio_h),thermal_height)
        box_temperature = thermal_row[y1:y2,x1:x2]
        if(box_temperature.size != 0):
            return box_temperature.max()
        else:
            return 0
if __name__ == '__main__':
    import sys
    import os
    app = QtWidgets.QApplication(sys.argv)
    window = Main()
    window.show()
    status = app.exec_()
    time.sleep(1)
#     sys.exit(app.exec_())
    os._exit(status)

