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


import cv2
import numpy as np

import time
import sys


def max_temperature(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

thermal = Thermal(width=160, height=120, framerate=9, frame_width=640, frame_height=480, log=None)
camera = Camera(width=640, height=480, framerate=30, log=None, channel=0, flip=1)
detector = FaceDetector(model_path="../Ultra-Light-Fast-Generic-Face-Detector-1MB/MNN/model/version-RFB/RFB-320.mnn" , input_size=(320,240))
anti = AntiSpoofing(model_path="../Ultra-Light-Fast-Generic-Face-Detector-1MB/MNN/model/4_0_0_80x80_MiniFASNetV1SE.mnn")
aligner = BoxAligner(img_h = 480,img_w = 640)
heater = Heater(pwm_pin=12 , target_temp=36.0)
correcter = Correcter(model_path="../thermal-tk/20210421_correcter.sav")

heater.start()
time.sleep(2)
thermal.start()
time.sleep(2)
camera.start()
time.sleep(2)

while(thermal.isOpened and camera.isOpened):
    thermal_ret,thermal_frame, thermal_row = thermal.capture()
    ret,frame = camera.capture()
    if(thermal_ret and ret):
        blackbody_max = thermal_row[:,:34].max()
        heater_temp = heater.last_temp
        boxes, labels, probs = detector.predict(frame)
        for i in range(boxes.shape[0]):
            box = boxes[i, :]
            score = probs[i]
            if(score > 0.9):
                #scaled_box = anti.scale_box(img_h=480,img_w=640,box=box)
                #expanded_face = frame[scaled_box[1]:scaled_box[3],scaled_box[0]:scaled_box[2]].copy()
                #live_score = anti.predict(expanded_face)
                live_score = 0.9
                if(live_score > 0.9):
                    color = (0, 255, 0)
                else:
                    color = ( 0, 0, 255)
                thermal_box = aligner.box_aligment([box])[0]
                cv2.rectangle(thermal_frame, (thermal_box[0], thermal_box[1]), (thermal_box[2], thermal_box[3]), color, 2)
                cv2.rectangle(frame, (box[0], box[1]), (box[2], box[3]), color, 2)
#                 cv2.rectangle(thermal_frame, (box[0], box[1]), (box[2], box[3]), color, 2)
#                 max_temp = thermal_row[thermal_box[1]//4:thermal_box[3]//4,thermal_box[0]//4:thermal_box[2]//4]
                face_temp = max_temperature(thermal_box,thermal_row,black_h=32)
#                 if(max_temp.size != 0):
                corrected_face_temp = correcter.predict(x_detect_k= face_temp, x_black_k= blackbody_max, x_black_c= heater_temp)
                cv2.putText(frame, "{:.2f}->{:.2f}".format(face_temp*0.0092,corrected_face_temp), ((box[0]+box[2])//2, (box[1]+box[3])//2), cv2.FONT_HERSHEY_COMPLEX_SMALL,1, (0, 255, 255), 1, cv2.LINE_AA)
                cv2.putText(thermal_frame, "{:.2f}".format(heater_temp), (320, 64), cv2.FONT_HERSHEY_COMPLEX_SMALL,1, (0, 0, 0), 1, cv2.LINE_AA)
        cv2.imshow('frames',np.column_stack((frame,thermal_frame)))
        if(cv2.waitKey(1) == ord('q')):
            break
    else:
        break
cv2.destroyAllWindows()

thermal.stop()
time.sleep(2)
camera.stop()
time.sleep(2)
heater.stop()
sys.exit()
