# 示意代码，不可直接运行：阈值参数与实现细节为示意，论文未公布完整代码
import numpy as np
from scipy.signal import welch

class RatioBasedThresholding:
    """
    比值阈值标定（RBT）
    
    论文核心方法：
    1. 计算清醒基线θ/β比值
    2. 动态计算当前θ/β比值
    3. 基于基线偏差生成疲劳标签
    """
    
    def __init__(self, fs: int = 200, 
                 theta_band=(4, 8),
                 beta_band=(13, 30)):
        self.fs = fs
        self.theta_band = theta_band
        self.beta_band = beta_band
        
    def compute_theta_beta_ratio(self, eeg: np.ndarray) -> float:
        """
        计算θ/β功率比值
        
        Args:
            eeg: [n_channels, n_samples] EEG信号
            
        Returns:
            ratio: θ/β功率比值
        """
        # 平均参考
        eeg = eeg - eeg.mean(axis=0)
        
        # Welch功率谱
        freqs, psd = welch(eeg, fs=self.fs, nperseg=512)
        
        # θ和β频段功率
        theta_mask = (freqs >= self.theta_band[0]) & \
                      (freqs < self.theta_band[1])
        beta_mask = (freqs >= self.beta_band[0]) & \
                     (freqs < self.beta_band[1])
        
        theta_power = np.sum(psd[:, theta_mask], axis=1).mean()
        beta_power = np.sum(psd[:, beta_mask], axis=1).mean()
        
        return theta_power / (beta_power + 1e-8)
    
    def calibrate_baseline(self, alert_eeg: np.ndarray, 
                            window_sec: int = 10) -> float:
        """
        从清醒状态EEG标定个体基线
        
        Args:
            alert_eeg: 清醒状态EEG数据
            window_sec: 窗口长度
            
        Returns:
            baseline_ratio: 基线θ/β比值
        """
        window_samples = window_sec * self.fs
        n_windows = len(alert_eeg[0]) // window_samples
        
        ratios = []
        for i in range(n_windows):
            window = alert_eeg[:, i*window_samples:(i+1)*window_samples]
            ratios.append(self.compute_theta_beta_ratio(window))
        
        baseline = np.mean(ratios)
        baseline_std = np.std(ratios)
        
        return baseline, baseline_std
    
    def generate_labels(self, eeg: np.ndarray,
                         baseline: float,
                         baseline_std: float,
                         threshold_std: float = 1.5) -> np.ndarray:
        """
        基于RBT生成疲劳标签
        
        Args:
            eeg: [n_channels, n_samples] 
            baseline: 基线θ/β比值
            baseline_std: 基线标准差
            threshold_std: 阈值标准差倍数
            
        Returns:
            labels: [n_windows] 0=清醒, 1=疲劳
        """
        window_samples = 10 * self.fs  # 10秒窗口
        n_windows = len(eeg[0]) // window_samples
        
        labels = np.zeros(n_windows)
        ratios = []
        
        for i in range(n_windows):
            window = eeg[:, i*window_samples:(i+1)*window_samples]
            ratio = self.compute_theta_beta_ratio(window)
            ratios.append(ratio)
            
            # RBT判定：超过基线+1.5σ为疲劳
            if ratio > baseline + threshold_std * baseline_std:
                labels[i] = 1
        
        return labels, ratios


# 测试
if __name__ == "__main__":
    rbt = RatioBasedThresholding(fs=200)
    
    # 模拟清醒EEG（低θ高β）
    np.random.seed(42)
    alert_eeg = np.random.randn(4, 200 * 60)  # 4通道60秒
    alert_eeg += 2.0 * np.sin(2 * np.pi * 20 * 
                                np.arange(200*60) / 200)  # 20Hz β
    
    # 标定基线
    baseline, baseline_std = rbt.calibrate_baseline(alert_eeg)
    print(f"基线θ/β比值: {baseline:.4f} ± {baseline_std:.4f}")
    
    # 模拟疲劳EEG（高θ低β）
    fatigue_eeg = np.random.randn(4, 200 * 60)
    fatigue_eeg += 3.0 * np.sin(2 * np.pi * 6 * 
                                 np.arange(200*60) / 200)  # 6Hz θ
    
    # 生成标签
    labels, ratios = rbt.generate_labels(
        fatigue_eeg, baseline, baseline_std
    )
    print(f"疲劳窗口比例: {labels.mean():.1%}")
    print(f"平均θ/β比值: {np.mean(ratios):.4f}")
    print(f"基线+1.5σ阈值: {baseline + 1.5 * baseline_std:.4f}")