bci_algo/concentration/algorithm/calculate_focus.py

import numpy as np
from scipy.signal import welch
from scipy.fft import fft
from scipy import signal
from collections import deque
import time
import os
# import logging
import base64
import io
import math

# logger = logging.getLogger(__name__)
#
# try:
#     import matplotlib
#     matplotlib.use('Agg')
#     import matplotlib.pyplot as plt
#     MATPLOTLIB_AVAILABLE = True
# except ImportError:
#     MATPLOTLIB_AVAILABLE = False
#     logger.warning("matplotlib未安装，报告图表功能不可用")


class Calculate():
    def __init__(self, Threshold_value_low, Threshold_value_high, fs=250, win_len=10, config=None):
        self.Threshold_value_low = Threshold_value_low
        self.Threshold_value_high = Threshold_value_high
        self.fs = fs
        self.focus_result = []
        self.CLI_result = []
        self.EVI_result = []
        self.eegQueue = deque(maxlen=win_len)

        # 初始化滤波器
        self.b_notch, self.a_notch = signal.iirnotch(50 / (self.fs/2), 30)
        self.b_design = signal.firwin(65, [2 / (self.fs/2), 40 / (self.fs/2)], pass_zero=False)

        self.last_focus = None
        # 异步滤波系数配置（核心手感控制纽）
        self.alpha_up = 1  # 上升系数：较小，保证分数平滑爬升，过滤偶发的瞬时高能量
        # alpha_down / shrink_factor 从 config.ini 读取，方便上位机调参
        if config:
            self.alpha_down = float(config.get('alpha_down', 0.8))
            self.shrink_factor = float(config.get('shrink_factor', 0.5))
        else:
            self.alpha_down = 0.8
            self.shrink_factor = 0.5
        print("[调试] Calculate 类初始化完成")

    def calculate_focus(self, beta, alpha, theta):
        """
        专注度计算 - 三区间门限异步滤波版本
        """
        # 0. 频带特征预处理
        theta_mod = theta ** 0.7

        # 原始比值
        raw = beta / (alpha + theta_mod + 1e-10)

        exponent = 2.0

        # 1. 防止脑电比值出现负数异常值
        raw_input = max(raw, 0.0)

        # 2. 2次幂纵轴压缩映射 (shrink_factor 从 config.ini 读取)
        focus_raw = 100 * self.shrink_factor * (raw_input ** exponent)

        # 3. 计算当前帧的瞬时分数 (基准量级 0-120)
        instant_focus = 120 * (1.0 - np.exp(-focus_raw / 100.0))

        # 4. 核心修改：三区间门限时域滤波
        if self.last_focus is None:
            # 冷启动：首帧直接赋值
            focus = instant_focus
        else:
            # 判断当前瞬时分数是否处于【极端区】(80以上 或 60以下)
            if instant_focus > 85.0 or instant_focus < 60.0:
                # 执行异步低通时域滤波
                if instant_focus >= self.last_focus:
                    # 趋势上升：慢爬升
                    focus = self.alpha_up * instant_focus + (1 - self.alpha_up) * self.last_focus
                else:
                    # 趋势下降：快跌落
                    focus = self.alpha_down * instant_focus + (1 - self.alpha_down) * self.last_focus
            else:
                # 【高灵敏自由区】(60 <= instant_focus <= 80)
                # 不执行异步滤波，分数直接跟随瞬时值，保证中间状态绝对跟手
                focus = instant_focus

        # 5. 更新历史状态缓存
        self.last_focus = focus

        # 打印在线调试日志，方便观察区间切换
        zone_tag = "极端区(滤波)" if (instant_focus > 80 or instant_focus < 60) else "自由区(直通)"
        print(f"原始特征比值 raw: {raw:.4f} | 瞬时分数: {instant_focus:.1f} | 滤波后分数: {focus:.1f}")

        # 最终返回整型
        return int(focus)


    def calculate_all(self, data, fs, nperseg=1000):
        mean_x = np.mean(data, axis=-1, keepdims=True)
        data = data - mean_x
        freqs, psd = self.compute_psd_multichannel(data, fs, nperseg)
        beta_psd = np.sum(self.band_psd(freqs, psd, (13, 30)))
        alpha_psd = np.sum(self.band_psd(freqs, psd, (8, 13)))
        theta_psd = np.sum(self.band_psd(freqs, psd, (4, 8)))
        
        print(f"[功率] β={beta_psd:.2f} | α={alpha_psd:.2f} | θ={theta_psd:.2f}")
        
        
        focus_score = self.calculate_focus(beta_psd, alpha_psd, theta_psd)
        focus_score = max(0, min(100, focus_score))
        
        self.focus_result.append(focus_score)
        if len(self.focus_result) > 3:
            self.focus_result.pop(0)
        final_focus = int(self.simple_moving_average(self.focus_result, window_size=5))

        cli_denom = alpha_psd + beta_psd
        CLI_score = np.log(theta_psd / (cli_denom + 1e-10)) if cli_denom > 0 else 0
        self.CLI_result.append(CLI_score)
        if len(self.CLI_result) > 5:
            self.CLI_result.pop(0)
        final_CLI = round(self.simple_moving_average(self.CLI_result, window_size=5), 2)
        
        return final_focus, final_CLI, beta_psd, alpha_psd, theta_psd

    def compute_psd_multichannel(self, data, fs=250, nperseg=1000):
        n_samples = data.shape[-1]
        if n_samples < nperseg:
            nperseg = n_samples
        
        noverlap = 500
        if noverlap >= nperseg:
            noverlap = int(nperseg / 2)

        if nperseg == 0:
            return np.array([]), np.zeros((data.shape[0], 0))

        freqs, psd = welch(data, fs=fs, nperseg=nperseg, noverlap=noverlap, axis=-1)
        return freqs, psd

    def band_psd(self, freqs, psd, band):
        idx = np.logical_and(freqs >= band[0], freqs <= band[1])
        return np.sum(psd[:, idx], axis=-1)

    def simple_moving_average(self, data, window_size=5):
        if len(data) == 0:
            return 30
        window = data[-window_size:]
        return sum(window) / len(window)

    def reset_queue(self):
        self.eegQueue.clear()

    # def start_recording(self):
    #     """开始记录数据"""
    #     self.recording = True
    #     self.start_time = time.time()
    #     self.beta_history = []
    #     self.alpha_history = []
    #     self.theta_history = []
    #     self.focus_history = []
    #     self.timestamp_history = []
    #     print("[调试] ========== 开始记录专注度数据 ==========")

    # def stop_recording(self):
    #     """停止记录并生成图表"""
    #     print(f"[调试] stop_recording被调用, recording={self.recording}, focus_history长度={len(self.focus_history)}")
    #     self.recording = False
    #     if len(self.focus_history) > 0:
    #         print("[调试] 数据非空，开始生成图表...")
    #         # 保存到本地文件
    #         chart_path = self.save_chart_to_file()
    #         if chart_path:
    #             print(f"[调试] 本地文件保存成功: {chart_path}")
    #         else:
    #             print("[调试] 本地文件保存失败")
    #         # 生成base64编码
    #         base64_data = self.generate_chart_base64()
    #         return base64_data
    #     else:
    #         print("[调试] 没有数据可保存，focus_history为空")
    #         return None

    # def add_data_point(self, focus, beta, alpha, theta):
    #     if not self.recording:
    #         return
    #     current_time = time.time()
    #     elapsed = current_time - self.start_time
    #
    #     self.beta_history.append(beta)
    #     self.alpha_history.append(alpha)
    #     self.theta_history.append(theta)
    #     self.focus_history.append(focus)
    #     self.timestamp_history.append(elapsed)
    #     print(f"[调试] 记录数据点: time={elapsed:.1f}s, focus={focus}, beta={beta:.2f}")

    # def save_chart_to_file(self):
    #     """
    #     保存图表到本地文件（唯一实现）
    #     """
    #     print(f"[调试] save_chart_to_file被调用, MATPLOTLIB_AVAILABLE={MATPLOTLIB_AVAILABLE}")
    #
    #     if not MATPLOTLIB_AVAILABLE:
    #         print("[调试] matplotlib不可用，无法保存")
    #         return None
    #
    #     if len(self.focus_history) < 2:
    #         print(f"[调试] 数据点不足，需要至少2个点，当前{len(self.focus_history)}个点")
    #         return None
    #
    #     print(f"[调试] 开始保存图表到本地文件...")
    #
    #     # 确保所有列表长度一致
    #     min_len = min(len(self.beta_history), len(self.alpha_history),
    #                   len(self.theta_history), len(self.focus_history),
    #                   len(self.timestamp_history))
    #
    #     print(f"[调试] 数据长度: min_len={min_len}")
    #
    #     beta_list = self.beta_history[:min_len]
    #     alpha_list = self.alpha_history[:min_len]
    #     theta_list = self.theta_history[:min_len]
    #     focus_list = self.focus_history[:min_len]
    #     times = self.timestamp_history[:min_len]
    #
    #     # 生成文件名
    #     timestamp = time.strftime("%Y%m%d_%H%M%S")
    #     chart_path = os.path.join(self.chart_dir, f"concentration_report_{timestamp}.png")
    #     print(f"[调试] 保存路径: {chart_path}")
    #
    #     try:
    #         # 创建图表
    #         fig, ax1 = plt.subplots(figsize=(14, 8))
    #
    #         # 左Y轴：功率数据
    #         ax1.plot(times, beta_list, 'b-', linewidth=1.5, alpha=0.8, label='Beta Power')
    #         ax1.plot(times, alpha_list, 'g-', linewidth=1.5, alpha=0.8, label='Alpha Power')
    #         ax1.plot(times, theta_list, 'orange', linewidth=1.5, alpha=0.8, label='Theta Power')
    #         ax1.set_xlabel('Time (s)', fontsize=12)
    #         ax1.set_ylabel('Band Power', fontsize=12, color='black')
    #         ax1.tick_params(axis='y', labelcolor='black')
    #         ax1.legend(loc='upper left')
    #         ax1.grid(True, alpha=0.3)
    #
    #         # 右Y轴：专注度
    #         ax2 = ax1.twinx()
    #         ax2.plot(times, focus_list, 'r-', linewidth=2, alpha=0.9, label='Focus (%)')
    #         ax2.set_ylabel('Focus (%)', fontsize=12, color='red')
    #         ax2.tick_params(axis='y', labelcolor='red')
    #         ax2.set_ylim(0, 105)
    #         ax2.legend(loc='upper right')
    #
    #         # 标题
    #         duration = times[-1] if times else 0
    #         avg_focus = np.mean(focus_list) if focus_list else 0
    #         plt.title(f'Concentration and EEG Band Power Trend\nDuration: {duration:.1f}s, Avg Focus: {avg_focus:.1f}%',
    #                   fontsize=14)
    #
    #         plt.tight_layout()
    #         plt.savefig(chart_path, dpi=150, bbox_inches='tight')
    #         plt.close()
    #
    #         print(f"\n{'='*60}")
    #         print(f"专注度报告图片已保存到本地:")
    #         print(f"  文件路径: {chart_path}")
    #         print(f"  数据点数: {min_len}")
    #         print(f"  时长: {duration:.1f}秒")
    #         print(f"  平均专注度: {avg_focus:.1f}%")
    #         print(f"{'='*60}\n")
    #
    #         return chart_path
    #
    #     except Exception as e:
    #         print(f"[调试] 保存文件时出错: {e}")
    #         import traceback
    #         traceback.print_exc()
    #         return None
    #
    # def generate_chart_base64(self):
    #     """
    #     生成图表的base64编码（用于网络传输）
    #     """
    #     if not MATPLOTLIB_AVAILABLE:
    #         return None
    #
    #     if len(self.focus_history) < 2:
    #         return None
    #
    #     min_len = min(len(self.beta_history), len(self.alpha_history),
    #                   len(self.theta_history), len(self.focus_history),
    #                   len(self.timestamp_history))
    #
    #     beta_list = self.beta_history[:min_len]
    #     alpha_list = self.alpha_history[:min_len]
    #     theta_list = self.theta_history[:min_len]
    #     focus_list = self.focus_history[:min_len]
    #     times = self.timestamp_history[:min_len]
    #
    #     fig, ax1 = plt.subplots(figsize=(14, 8))
    #
    #     ax1.plot(times, beta_list, 'b-', linewidth=1.5, alpha=0.8, label='Beta Power')
    #     ax1.plot(times, alpha_list, 'g-', linewidth=1.5, alpha=0.8, label='Alpha Power')
    #     ax1.plot(times, theta_list, 'orange', linewidth=1.5, alpha=0.8, label='Theta Power')
    #     ax1.set_xlabel('Time (s)', fontsize=12)
    #     ax1.set_ylabel('Band Power', fontsize=12, color='black')
    #     ax1.tick_params(axis='y', labelcolor='black')
    #     ax1.legend(loc='upper left')
    #     ax1.grid(True, alpha=0.3)
    #
    #     ax2 = ax1.twinx()
    #     ax2.plot(times, focus_list, 'r-', linewidth=2, alpha=0.9, label='Focus (%)')
    #     ax2.set_ylabel('Focus (%)', fontsize=12, color='red')
    #     ax2.tick_params(axis='y', labelcolor='red')
    #     ax2.set_ylim(0, 105)
    #     ax2.legend(loc='upper right')
    #
    #     duration = times[-1] if times else 0
    #     avg_focus = np.mean(focus_list) if focus_list else 0
    #     plt.title(f'Concentration and EEG Band Power Trend\nDuration: {duration:.1f}s, Avg Focus: {avg_focus:.1f}%',
    #               fontsize=14)
    #
    #     plt.tight_layout()
    #
    #     buffer = io.BytesIO()
    #     plt.savefig(buffer, format='png', dpi=150, bbox_inches='tight')
    #     buffer.seek(0)
    #     image_base64 = base64.b64encode(buffer.read()).decode('utf-8')
    #     plt.close()
    #
    #     return image_base64

    def queueOpt(self, data):
        if data is None or data.size == 0:
            return None
        if len(self.eegQueue) < self.eegQueue.maxlen:
            self.eegQueue.append(data)
        else:
            self.eegQueue.append(data)

        if len(self.eegQueue) == self.eegQueue.maxlen:
            eegData = np.hstack([self.eegQueue[i] for i in range(len(self.eegQueue))])
            if eegData.size == 0:
                return None
            eegData -= np.mean(eegData, axis=-1, keepdims=True)
            # eegData = signal.lfilter(self.b_notch, self.a_notch, eegData)    # 陷波
            # eegData = signal.lfilter(self.b_design, 1, eegData)              # 滤波
            focus_score, CLI_score, beta_psd, alpha_psd, theta_psd = self.calculate_all(eegData, fs=self.fs, nperseg=1000)

            # self.add_data_point(focus_score, beta_psd, alpha_psd, theta_psd)  # 已注释（方法已移除）

            # return (focus_score)
            return (focus_score, beta_psd)
        # return None



class Calculate2():
    def __init__(self, Threshold_value_low, Threshold_value_high):
        self.Threshold_value_low = Threshold_value_low
        self.Threshold_value_high = Threshold_value_high
        self.focus_result = []
        self.theta_result = []
        self.alpha_result = []
        self.flow_result = []

    def calculate_all(self, data, fs, L=2500):
        mean_x = np.mean(data, axis=-1, keepdims=True)
        data = data - mean_x

        Y = fft(data, axis=-1)
        P2 = np.abs(Y / L)
        P1 = P2[:, :L // 2 + 1]
        P1[:, 1:-1] = 2 * P1[:, 1:-1]

        beta_power = self.PSD(P1, L, fs, 13, 30)
        alpha_power = self.PSD(P1, L, fs, 8, 13)
        theta_power = self.PSD(P1, L, fs, 4, 8)
        gamma_power = self.PSD(P1, L, fs, 30, 100)
        
        focus_score = beta_power / (alpha_power + theta_power)
        print('focus score:', focus_score)
        focus_score = ((focus_score - self.Threshold_value_low) * 100) / (self.Threshold_value_high - self.Threshold_value_low)
        self.focus_result.append(focus_score)
        if len(self.focus_result) > 3:
            self.focus_result.pop(0)
        final_focus = int(self.simple_moving_average(self.focus_result, window_size=3))
        
        self.theta_result.append(theta_power)
        if len(self.theta_result) > 30:
            self.theta_result.pop(0)
        self.alpha_result.append(alpha_power)
        if len(self.alpha_result) > 30:
            self.alpha_result.pop(0)
        rest_theta = self.simple_moving_average(self.theta_result, window_size=30)
        rest_alpha = self.simple_moving_average(self.alpha_result, window_size=30)
        distraction_score = (theta_power / rest_theta) * (1 - (alpha_power / rest_alpha))
        
        flow_score = gamma_power / beta_power
        flow_score = (flow_score / self.Threshold_value_high) * 100
        self.flow_result.append(flow_score)
        if len(self.flow_result) > 3:
            self.flow_result.pop(0)
        final_flow = int(self.simple_moving_average(self.flow_result, window_size=3))

        return final_focus, distraction_score, final_flow

    def PSD(self, P1, L, Fs, s_freq, e_freq):
        s_point = round(s_freq * L / Fs)
        e_point = round(e_freq * L / Fs)
        x, y = P1.shape
        band_PSD = 0
        for i in range(x):
            for j in range(s_point, e_point):
                band_PSD += P1[i, j] ** 2
        return band_PSD

    def simple_moving_average(self, data, window_size=3):
        if len(data) == 0:
            return []
        window = data[-window_size:]
        return sum(window) / len(window)