import queue

import numpy as np
from PyQt5.QtCore import QObject, pyqtSignal
from multiprocessing import Process, Queue
from apis.age.AgeGenderPredictor import AgeGenderPredictor
from apis.hr.HeartRateMonitor import HeartRateMonitor
from apis.st.predict_api import SkinTypePredictor
from apis.sd.predict_api import SkinDiseasePredictor
from apis.emotion.predict_api import EmotionPredictor
from apis.bp.BPApi import BPModel
from apis.rr.RespirationRateDetector import RespirationRateDetector
from apis.rr.params import args


def run_ai(input_queue, output_queue, fps_queue):
    age_predictor = AgeGenderPredictor('weights/age.pth')
    st_predictor = SkinTypePredictor('weights/st.pth', class_indices_path="labels/st.json")
    sd_predictor = SkinDiseasePredictor('weights/sd.pth', class_indices_path="labels/sd.json")
    emotion_predictor = EmotionPredictor('weights/emotion.pth', class_indices_path="labels/emotion.json")
    bp_predictor = BPModel(model_path=r'weights/bp.pth', fps=30)
    rr_detector = RespirationRateDetector(args)
    hr_detector = HeartRateMonitor(30, 0.8, 1.8)

    while True:

        frames = input_queue.get()
        fps = fps_queue.get()

        if frames is None:
            break

        if fps is None:
            fps = 30
        else:
            # print(fps)
            bp_predictor.fps = fps
            hr_detector.fps = fps

        results = {}

        sbp_outputs, dbp_outputs = bp_predictor.predict(frames[-300:])

        gender, age, age_group = age_predictor.predict(frames[-1])
        st = st_predictor.predict(frames[-1])
        sd = sd_predictor.predict(frames[-1])
        emotion = emotion_predictor.predict(frames[-1])

        Resp, RR_FFT, RR_PC, RR_CP, RR_NFCP = rr_detector.detect_respiration_rate(frames[-300:], fps)
        heartrate, sdnn, rmssd, cv_rr, bvp = hr_detector.process_roi(frames[-300:])

        RR = int(round((RR_FFT + RR_PC + RR_CP + RR_NFCP) / 4))

        results['gender'] = gender
        results['age'] = int(age)
        results['age_group'] = age_group
        results['skin_type'] = st
        results['skin_disease'] = sd
        results['emotion'] = emotion
        results['sbp'] = int(round(sbp_outputs, 0))
        results['dbp'] = int(round(dbp_outputs, 0))
        results['rr'] = RR
        results['hr'] = int(heartrate)
        results['sdnn'] = round(sdnn, 1)

        # 判断rmssd是否为nan
        results['rmssd'] = round(rmssd, 1) if not np.isnan(rmssd) else 0.0
        results['cvrr'] = round(cv_rr, 1)
        results['resp'] = Resp
        results['bvp'] = bvp

        # 添加其他生理指标的计算...

        def clear_queue(q):
            while not q.empty():
                try:
                    q.get_nowait()
                except queue.Empty:
                    break

        clear_queue(input_queue)
        output_queue.put(results)


class AIProcess(QObject):
    results_ready = pyqtSignal(dict)

    def __init__(self):
        super().__init__()
        self.input_queue = Queue()
        self.output_queue = Queue()
        self.fps_queue = Queue()

        self.process = Process(target=run_ai, args=(self.input_queue, self.output_queue, self.fps_queue))

    def start(self):
        self.process.start()

    def process_frames(self, frames):
        self.input_queue.put(frames)

    def update_fps(self, fps):
        self.fps_queue.put(fps)

    def stop(self):
        self.input_queue.put(None)
        self.process.join()

    def get_results(self):
        if not self.output_queue.empty():
            return self.output_queue.get()
        return None