import ast
import threading
from datetime import datetime
import multiprocessing as mp
import numpy as np
import time
import torch
from queue import Empty
from scipy import signal
from torch.autograd import Variable
from Device.SunnyLinker import SunnyLinker64
from MI.Algorithm.otherModels import weights_init
from SSMVEP.algorithm.tdca import TDCA
from SSMVEP.algorithm.base import generate_cca_references
from Zmq.zmqServer import zmqServer
from Zmq.zmqClient import zmqClient
from MI.Algorithm.conformer_2class import onlineTrain
from PubLibrary.InifileHelper import IniRead
from SSVEP.dwfbcca import FbccaDw
from Tools.plot_MI_EEG import plotMain

class Decoder_main(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.Runing=True
        self.decoder = None

        self.fs = 250  # 采样率
        self.energy = 0  # 电量
        self.status_code = 0  # 与采集设备通信的状态码，0为异常，1为正常
        self.decoder_class = None #解码器类别

        self.decodingSteps = 0  # 0=停止解码 1=预热 2=解码中 3=解码完成，发送解码结果
    def connect(self):
        self.thread_data_server = SunnyLinker64('127.0.0.1', 7878, 250, 64,
                                                method='tcp')
        self.thread_data_server.toUv = True
        self.thread_data_server.start()

        self.zmqServer = zmqServer()
        self.zmqServer.start()
        self.zmqClient = zmqClient('127.0.0.1', 8088)
        self.zmqClient.connect()

    def is_valid_signal(self, data, threshold=1e5):  # 判断当前信号是否为有效信号
        # data: (chans, samples)
        energy = np.mean(np.var(data, axis=1))  # 各通道方差均值
        if energy > threshold:
            return False
        return True
    def init_Decoder(self,decoder_class):
        '''
        初始化解码器
        :param decoder_class: 'ssvep' or 'ssmvep' or 'mi'
        :return:
        '''
        self.decoder_class = decoder_class
        if decoder_class == 'ssvep':
            self.n_chan = 8
            self.thread_data_server.interval_inited = False
            DW_cost_method, self.DW_cost_tv = ast.literal_eval(IniRead('system', 'SSVEP_ThresholdValue'))
            self.ListFreq = self.zmqServer.targetFreqs
            self.num_target = len(self.ListFreq)
            if self.num_target == 0:
                return
            # 初始化对象 二代算法
            self.dw = FbccaDw(self.fs, self.num_target, self.n_chan, 5, 5,
                              0.2, [2.0, 0.1], [8, 7], 50, DW_cost_method)
            #  frequence band
            self.dw.filterFrequenceBank()
            self.dw.setNotchFilterPara()
            self.calculateCount = 0
            self.referenceData = self.dw.reference(self.ListFreq, int(50 * 0.2 * self.fs),
                                                   5)
            self.dw.filterInit()
            self.dw.onlineInit()  # 刺激闪烁的第1s重置 --在线数据采集时

        elif decoder_class == 'ssmvep':
            self.thread_data_server.interval_init(decoder_class)
            self.n_chan = 8
            self.interval_epoch = ast.literal_eval(IniRead('system', 'SSMVEP_IntervalEpoch'))
            self.sample_length = round(self.interval_epoch[1] - self.interval_epoch[0], 6)  # 解码数据长度2s,# 精确到小数点后6位
            self.single_train = 10  # 单类别数量
            self.num_target = 2  # 分类目标数目
            self.list_freqs = np.array([8, 9])  # 刺激频率
            self.list_phase = np.array([0, 0])  # 相位
            self.tdca = TDCA(padding_len=5, n_components=1)
            self.Yf = generate_cca_references(self.list_freqs, srate=self.fs, T=self.sample_length,
                                              phases=self.list_phase, n_harmonics=5)
            self.parameter_init(5,45)
        elif decoder_class == 'mi':
            self.thread_data_server.interval_init(decoder_class)
            self.n_chan = 21
            self.interval_epoch = ast.literal_eval(IniRead('system', 'MI_IntervalEpoch'))
            self.sample_length = round(self.interval_epoch[1] - self.interval_epoch[0], 6)  # 解码数据长度2s,# 精确到小数点后6位
            self.single_train = 40  # 单类别数量
            self.num_target = 2  # 分类目标数目

            self.parameter_init(8, 30)


    def parameter_init(self,bandPass_low,bandPass_high):
        self.interval_epoch = [int(i * self.fs) for i in self.interval_epoch]  # epoch截取信息
        self.train_epoch = [int(self.interval_epoch[0]), int(self.interval_epoch[1] + 0.1 * self.fs)]  # 训练样本epoch
        self.trainData = [] #训练数据
        self.trainLabel = [] #训练标签
        self.plotData = [] #报告分析数据
        self.plotLabel = [] #报告分析标签
        self.currentLabel = -1 #刺激界面当前显示的训练标签
        self.train_started = False #是否开始训练模型
        self.load_model = False  # 调用模型是否完成的标志
        self.b_notch, self.a_notch = signal.iirnotch(50 / (self.fs/2), 30)  # 50Hz工频陷波，250是采样率，30是质量因子
        self.b_design = signal.firwin(65, [bandPass_low / (self.fs/2), bandPass_high / (self.fs/2)], pass_zero=False)  # 设计8-30Hz带通滤波器
        fileName = 'Model_' + datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        filePath = './online_Models/'
        self.modelPath = ''.join([filePath, fileName, '.pth'])
        self.mp_data_queue = mp.Queue() #多进程传参队列
        self.mp_result_queue = mp.Queue() #多进程结果队列

    def preprocess(self, signal_data):
        # # 计算每行的平均值
        row_means = np.mean(signal_data, axis=-1, keepdims=True)
        # 对每一行去均值
        signal_data = signal_data - row_means

        signal_data = signal.lfilter(self.b_notch, self.a_notch, signal_data, axis=-1)  # 工频陷波
        signal_data = signal.lfilter(self.b_design, 1, signal_data, axis=-1)  # 带通滤波
        return signal_data

    def run(self):
        while self.Runing:
            if self.zmqServer.decoder_switch or self.zmqServer.changeTarget:
                self.zmqServer.decoder_switch = False
                self.zmqServer.changeTarget = False
                self.init_Decoder(self.zmqServer.decoder_class)

            # 同步信息
            if self.zmqServer.state_mode == 'sync':
                self.zmqClient.send_to_all('sync', self.zmqClient.state)
                self.zmqServer.state_mode = 'rest'
            # 状态异常，报告上位机
            if self.status_code != self.thread_data_server.status_code:
                self.status_code = self.thread_data_server.status_code
                self.zmqClient.send_to_all('status_code', int(self.status_code))

            # 返回电量
            if self.energy != self.thread_data_server.energy:
                self.energy = self.thread_data_server.energy
                self.zmqClient.send_to_all('energy', int(self.energy))

            if self.zmqServer.open_Impedance == True:  # 开启阻抗检测功能，仅运行一次
                self.thread_data_server.Impedance(True)
                self.zmqServer.open_Impedance = -1
            elif self.zmqServer.open_Impedance == False:
                self.thread_data_server.Impedance(False)
                self.zmqServer.open_Impedance = -1

            if self.zmqServer.get_Impedance:  # 返回阻抗值
                if self.thread_data_server.GetDataLenCount() > 250:
                    Impe_data = self.thread_data_server.getData(250)
                    # 计算阻抗
                    imps = self.thread_data_server.getImpedance(Impe_data,self.zmqServer.decoder_class)
                    self.zmqClient.send_to_all('impedance', imps.tolist())
                else:
                    pass
            if self.zmqServer.getReport:  #返回训练报告内容
                self.zmqServer.getReport = False
                allData = np.array(self.plotData)
                allLabel = np.array(self.plotLabel) + 1
                nTrials = min(len(allLabel),len(allData))
                if nTrials == 0:
                    self.zmqClient.send_to_all('miReport',0)
                else:
                    allData = allData[:nTrials]
                    allLabel = allLabel[:nTrials]
                    ch_names = ['FC3', 'FC1', 'FCZ', 'FC2', 'FC4', 'C5', 'C3', 'C1', 'CZ', 'C2', 'C4', 'C6', 'CP3', 'CP1',
                                'CP2', 'CP4', 'P3', 'P1', 'PZ', 'P2', 'P4']
                    compare_names = ['C3', 'CZ', 'C4']
                    miReport = plotMain(ch_names=ch_names,compare_names=compare_names,Data=allData,labels=allLabel,MI_label=1,Rest_label=2,
                             fs=self.fs)
                    self.zmqClient.send_to_all('miReport',miReport)


            if self.decoder_class == 'ssvep':
                self.decoder_SSVEP()
            elif self.decoder_class == 'ssmvep':
                self.decoder_SSMVEP()
            elif self.decoder_class == 'mi':
                self.decoder_MI()
            else:
                if self.zmqServer.get_Impedance == False:  # 非阻抗检测状态
                    if self.thread_data_server.GetDataLenCount() < 25:
                        time.sleep(0.005)
                        continue;
                    self.thread_data_server.getData(25)


    def decoder_SSVEP(self):
        if self.zmqServer.StartDecode:
            self.zmqServer.StartDecode = False
            self.decodingSteps = 1
            self.thread_data_server.ResetAll()
            print('启动预测')
        if self.thread_data_server.GetDataLenCount() < 50:
            time.sleep(0.005)
            return
        if self.zmqServer.get_Impedance == False:  # 非阻抗检测状态
            data = self.thread_data_server.getDataViaSSVEP(50)
            data = data[:self.n_chan, :]
        if self.decodingSteps == 1 and hasattr(self,'dw'):  # 开始预热
            self.dw.onlineInit()  # 刺激闪烁的第1s重置 --在线数据采集时
            self.dw.warmFilter(data)  # 预热
            self.decodingSteps = 2
            print('预热数据完成。开始预测')
            return
        if self.decodingSteps == 2 and hasattr(self,'dw'):  # 解码中
            choosenNum = self.dw.fbccaDWMW(data, self.referenceData, self.DW_cost_tv, self.calculateCount)
            self.calculateCount += 1
            if choosenNum != -1 and self.is_valid_signal(data):
                self.decodingSteps = 3
                print('预测结果：' + str(choosenNum) + ',计算次数：' + str(self.calculateCount))
                self.calculateCount = 0
        if self.decodingSteps == 3:  # 发送解码后的信息
            self.zmqClient.send_to_all('result', int(choosenNum))
            self.decodingSteps = 0
            print('发送给界面完成。')
    def decoder_SSMVEP(self):
        '''模型训练'''
        if self.load_model == False and all(
                self.trainLabel.count(i) >= self.single_train for i in range(len(self.list_freqs))):  # 模型尚未训练完成
            self.trainData = np.array(self.trainData)
            self.trainLabel = np.array(self.trainLabel)
            print(np.shape(self.trainData), (self.trainLabel))
            # 保存多个数组到文件
            # np.savez('20250520_yy.npz', array1=self.trainData, array2=self.trainLabel)
            # self.decoder = self.fbtdca.fit(self.trainData, self.trainLabel, Yf=self.Yf)
            self.decoder = self.tdca.fit(self.trainData, self.trainLabel, Yf=self.Yf)
            now = datetime.now()
            formatted_time = now.strftime('%H:%M:%S.%f')[:-3]
            print('模型训练完成', formatted_time)
            self.load_model = True
            self.zmqClient.send_to_all('paradigm', 1)

        '''训练阶段采集数据'''
        if self.zmqServer.state_mode == 'train':  # 训练状态
            if self.zmqServer.StartTrain:
                self.currentLabel = self.zmqServer.currentLabel
                self.zmqServer.StartTrain = False
            if self.thread_data_server.epoch_finished == False or self.thread_data_server.GetDataLenCount() < \
                    self.train_epoch[1] \
                    + self.thread_data_server.event_inner_idx:
                time.sleep(0.0001)
                return
            print('训练队列数据：', self.thread_data_server.GetDataLenCount())
            trainTrial = self.thread_data_server.get_SSMVEPData()  # 取出所有数据
            print('取出的： ', trainTrial.shape, 'event: ', trainTrial[-2, self.thread_data_server.event_inner_idx])
            trainTrial = self.preprocess(trainTrial[:self.n_chan, :])  # 预处理
            trainTrial = trainTrial[:, self.thread_data_server.event_inner_idx + self.train_epoch[
                0]:self.thread_data_server.event_inner_idx + self.train_epoch[1]]
            print('trial: ', self.thread_data_server.event_inner_idx, self.train_epoch[0], self.train_epoch[1])
            if trainTrial.shape[1] == (self.train_epoch[1] - self.train_epoch[0]) and isinstance(
                    self.trainLabel, list) \
                    and self.trainLabel.count(self.currentLabel) < self.single_train:
                self.trainData.append(trainTrial)
                self.trainLabel.append(self.currentLabel)

        elif self.zmqServer.state_mode == 'predict':  # 测试状态
            if self.load_model == False:  # 模型尚未训练完成
                time.sleep(0.01)
                return
            else:  # 已有模型
                if self.zmqServer.StartDecode:
                    self.zmqServer.StartDecode = False
                    now = datetime.now()
                    formatted_time = now.strftime('%H:%M:%S.%f')[:-3]
                    print('启动预测  ', formatted_time)

                if self.thread_data_server.epoch_finished == False or self.thread_data_server.GetDataLenCount() < \
                        self.interval_epoch[1] \
                        + self.thread_data_server.event_inner_idx:
                    time.sleep(0.0001)
                    return
                data = self.thread_data_server.get_SSMVEPData()  # 读取全部数据
                print('取出的： ', data.shape, 'event: ', data[-2, self.thread_data_server.event_inner_idx])
                data = self.preprocess(data[:self.n_chan, :])  # 预处理
                data = data[:,
                       self.thread_data_server.event_inner_idx + self.interval_epoch[
                           0]:self.thread_data_server.event_inner_idx + self.interval_epoch[1]]
                pad_eeg_test = np.zeros(
                    (data.shape[0], int((self.sample_length + 0.1) * self.fs)))
                pad_eeg_test[:, :int(self.sample_length * self.fs)] = data
                choosenNum, features_2 = self.decoder.predict(pad_eeg_test)
                if isinstance(choosenNum, np.ndarray):
                    choosenNum = choosenNum[0]
                print('结果：', choosenNum, 'rho: ', sorted(features_2[0]),
                      sorted(features_2[0])[-1] - sorted(features_2[0])[-2])
                self.zmqClient.send_to_all('result', int(choosenNum))
                print('发送给界面完成。')
        else:  # 休息状态
            if self.zmqServer.get_Impedance == False:  # 非阻抗检测状态
                if self.thread_data_server.GetDataLenCount() < 25:
                    time.sleep(0.005)
                    return 
                self.thread_data_server.getData(25)
    def decoder_MI(self):
        '''模型训练'''
        if self.train_started == False and all(
                self.trainLabel.count(i) >= self.single_train for i in range(self.num_target)):  # 模型尚未训练
            self.zmqClient.send_to_all('paradigm', 2)  # 模型训练前，训练集采集完毕，通知上位机
            self.train_started = True
            self.trainData = np.array(self.trainData)
            self.trainLabel = np.array(self.trainLabel) + 1
            # print('训练集：',np.shape(self.trainData), (self.trainLabel))
            p = mp.Process(target=onlineTrain, args=(self.mp_data_queue, self.mp_result_queue))  # 开启子进程，训练模型
            p.start()
            self.mp_data_queue.put({'data': self.trainData, 'label': self.trainLabel, 'modelPath': self.modelPath,
                                    'n_chan': self.n_chan})

        '''检查模型是否训练完成,调用'''
        if self.load_model == False and self.train_started == True:
            try:
                result = self.mp_result_queue.get_nowait()
                if result['status'] == 'success':
                    print("模型训练完成，加载新模型")
                    # 调用模型
                    self.model = torch.load(self.modelPath, weights_only=False)
                    self.model.eval()
                    # 模型预热
                    warmup_data = np.random.uniform(-1, 1, (1, 1, self.n_chan, 1000))
                    warmup_data = torch.from_numpy(warmup_data)
                    warmup_data = Variable(warmup_data.type(torch.cuda.FloatTensor))
                    with torch.no_grad():
                        _ = self.model(warmup_data)
                    self.load_model = True
                    self.zmqClient.send_to_all('paradigm', 1)  # 模型调用完毕，通知上位机
                else:
                    print("训练失败:", result['msg'])
            except Empty:
                pass  # 还没完成
            except Exception as e:
                print('模型调用失败: ', e)

        '''训练阶段采集数据'''
        if self.zmqServer.state_mode == 'train' and self.train_started == False:  # 训练状态
            if self.zmqServer.StartTrain:
                self.currentLabel = self.zmqServer.currentLabel
                self.zmqServer.StartTrain = False
            if self.thread_data_server.epoch_finished == False or self.thread_data_server.GetDataLenCount() < \
                    self.interval_epoch[1] \
                    + self.thread_data_server.event_inner_idx:
                time.sleep(0.0001)
                return
            print('训练队列数据：', self.thread_data_server.GetDataLenCount())
            originalTrial = self.thread_data_server.get_MIData()  # 取出MI导联数据
            print('取出的： ', originalTrial.shape, 'event: ', originalTrial[-2, self.thread_data_server.event_inner_idx])
            trainTrial = self.preprocess(originalTrial[:self.n_chan, :])  # 预处理
            trainTrial = trainTrial[:, self.thread_data_server.event_inner_idx + self.interval_epoch[
                0]:self.thread_data_server.event_inner_idx + self.interval_epoch[1]]
            print('trial: ', self.thread_data_server.event_inner_idx, self.interval_epoch[0], self.interval_epoch[1])
            if trainTrial.shape[1] == (self.interval_epoch[1] - self.interval_epoch[0]) and isinstance(self.trainLabel,
                                                                                                       list) \
                    and self.trainLabel.count(self.currentLabel) < self.single_train:
                self.trainData.append(trainTrial)
                self.trainLabel.append(self.currentLabel)
                print('训练集：', np.shape(self.trainData))
                self.plotData.append(originalTrial[:self.n_chan, self.thread_data_server.event_inner_idx + self.interval_epoch[
                0]:self.thread_data_server.event_inner_idx + self.interval_epoch[1]])
                self.plotLabel.append(self.currentLabel)

        elif self.zmqServer.state_mode == 'predict' and self.load_model == True:  # 测试状态
            if self.zmqServer.StartDecode:
                self.zmqServer.StartDecode = False
                now = datetime.now()
                formatted_time = now.strftime('%H:%M:%S.%f')[:-3]
                print('启动预测  ', formatted_time)

            if self.thread_data_server.epoch_finished == False or self.thread_data_server.GetDataLenCount() < \
                    self.interval_epoch[1] \
                    + self.thread_data_server.event_inner_idx:
                time.sleep(0.0001)
                return
            originalData = self.thread_data_server.get_MIData()  # 读取全部数据
            print('取出的： ', originalData.shape, 'event: ', originalData[-2, self.thread_data_server.event_inner_idx])
            start = time.time()
            data = self.preprocess(originalData[:self.n_chan, :])  # 预处理
            data = data[:,
                   self.thread_data_server.event_inner_idx + self.interval_epoch[
                       0]:self.thread_data_server.event_inner_idx + self.interval_epoch[1]]
            self.plotData.append(
                originalData[:self.n_chan, self.thread_data_server.event_inner_idx + self.interval_epoch[
                    0]:self.thread_data_server.event_inner_idx + self.interval_epoch[1]])

            test_data = data[np.newaxis, np.newaxis, :, :]
            test_data = torch.from_numpy(test_data)
            test_data = Variable(test_data.type(torch.cuda.FloatTensor))
            with torch.no_grad():
                Cls = self.model(test_data)
                y_pred = torch.max(Cls, 1)[1]
                self.plotLabel.append(int(y_pred.item()))
                print('运动意图识别: ', y_pred)
            self.zmqClient.send_to_all('result', int(y_pred.item()))
            end = time.time()
            print(f'发送给界面完成,耗时{end - start:.3f}s。')
        else:  # 休息状态
            if self.zmqServer.get_Impedance == False:  # 非阻抗检测状态
                if self.thread_data_server.GetDataLenCount() < 25:
                    time.sleep(0.005)
                    return
                self.thread_data_server.getData(25)

    def stop(self):
        '''
        停止运行
        @return:
        '''
        self.zmqServer.stop()
        self.Runing=False