"""
Risk-aware Sequential Labeling (RSL) 实现

核心思想：
1. 不是每帧独立标注，而是考虑序列上下文
2. 引入"风险过渡"概念，建模从安全到危险的渐变过程
3. 多标注者投票 + 序列一致性约束
"""

import numpy as np
from dataclasses import dataclass
from typing import List, Tuple
from enum import IntEnum


class RiskLevel(IntEnum):
    SAFE = 0
    LOW_RISK = 1
    HIGH_RISK = 2
    DANGER = 3


@dataclass
class EEGSegment:
    """EEG 数据段"""
    data: np.ndarray        # shape=(channels, time_steps)
    sample_rate: int
    start_time: float
    end_time: float
    label: int = -1          # -1 = unlabeled


class RSLLabeler:
    """
    Risk-aware Sequential Labeling
    
    步骤：
    1. 多标注者独立标注
    2. 序列一致性检查
    3. 风险过渡平滑
    4. 多数投票 + 置信度估计
    """
    
    def __init__(self, num_annotators=5, 
                 smooth_window=5,
                 transition_penalty=0.3):
        self.num_annotators = num_annotators
        self.smooth_window = smooth_window
        self.transition_penalty = transition_penalty
    
    def label_sequence(self, 
                       segments: List[EEGSegment],
                       annotator_labels: np.ndarray) -> List[int]:
        """
        对 EEG 序列进行 RSL 标注
        
        Args:
            segments: EEG 段列表
            annotator_labels: shape=(num_annotators, num_segments)
                             每个标注者的标签
            
        Returns:
            final_labels: 最终标签序列
        """
        num_segments = len(segments)
        
        # Step 1: 多数投票
        vote_labels = []
        for i in range(num_segments):
            votes = annotator_labels[:, i]
            unique, counts = np.unique(votes, return_counts=True)
            label = unique[np.argmax(counts)]
            vote_labels.append(int(label))
        vote_labels = np.array(vote_labels)
        
        # Step 2: 序列一致性检查
        # 检测突变（单帧跳变后恢复）
        for i in range(1, num_segments - 1):
            if (vote_labels[i] != vote_labels[i-1] and 
                vote_labels[i] != vote_labels[i+1] and
                vote_labels[i-1] == vote_labels[i+1]):
                # 突变，可能是标注噪声
                vote_labels[i] = vote_labels[i-1]
        
        # Step 3: 风险过渡平滑
        # 使用滑动窗口加权平均
        smoothed = vote_labels.copy()
        half_w = self.smooth_window // 2
        for i in range(half_w, num_segments - half_w):
            window = vote_labels[i-half_w:i+half_w+1]
            # 加权：中心权重最大
            weights = np.exp(-0.5 * (
                np.arange(len(window)) - half_w
            )**2 / (half_w/2)**2)
            weighted_avg = np.average(window, weights=weights)
            # 只在过渡区域平滑
            if abs(weighted_avg - vote_labels[i]) < 0.5:
                smoothed[i] = int(round(weighted_avg))
        
        # Step 4: 过渡惩罚
        # 相邻标签变化越大，惩罚越大
        transitions = np.diff(smoothed)
        for i, t in enumerate(transitions):
            if abs(t) >= 2:
                # 大跳变（如 SAFE -> DANGER），插入中间标签
                if t > 0:
                    smoothed[i] = max(smoothed[i], 
                                     smoothed[i] - 1)
                else:
                    smoothed[i] = min(smoothed[i], 
                                     smoothed[i] + 1)
        
        return smoothed.tolist()


# ========== 测试 ==========

if __name__ == "__main__":
    np.random.seed(42)
    
    # 模拟 100 个 EEG 段
    num_segments = 100
    num_annotators = 5
    
    # 真实标签序列：前30帧 SAFE, 中间20帧 LOW_RISK,
    # 然后30帧 HIGH_RISK, 最后20帧 DANGER
    true_labels = (
        [RiskLevel.SAFE] * 30 +
        [RiskLevel.LOW_RISK] * 20 +
        [RiskLevel.HIGH_RISK] * 30 +
        [RiskLevel.DANGER] * 20
    )
    
    # 模拟标注者（各有 10-15% 噪声）
    annotator_labels = np.zeros(
        (num_annotators, num_segments), dtype=int
    )
    for a in range(num_annotators):
        noise_mask = np.random.random(num_segments) < 0.12
        noisy_labels = np.array(true_labels).copy()
        # 随机偏移一些标签
        noisy_labels[noise_mask] = np.random.randint(
            0, 4, size=noise_mask.sum()
        )
        annotator_labels[a] = noisy_labels
    
    # RSL 标注
    labeler = RSLLabeler(num_annotators=5, smooth_window=7)
    segments = [
        EEGSegment(
            data=np.random.randn(32, 100),
            sample_rate=250,
            start_time=i*0.4,
            end_time=(i+1)*0.4
        )
        for i in range(num_segments)
    ]
    
    final_labels = labeler.label_sequence(
        segments, annotator_labels
    )
    
    # 评估
    accuracy = np.mean(
        np.array(final_labels) == np.array(true_labels)
    )
    print(f"=== RSL 标注结果 ===")
    print(f"  段数: {num_segments}")
    print(f"  标注者数: {num_annotators}")
    print(f"  标注准确率: {accuracy*100:.1f}%")
    print(f"  真实标签序列: {true_labels[:10]}...")
    print(f"  RSL 标签序列:  {final_labels[:10]}...")

"""
3D-CRNN (3D Convolutional Recurrent Neural Network)
论文核心模型复现

架构：
1. 3D 卷积层：提取 EEG 时空特征
2. 循环层（LSTM/GRU）：建模时间动态
3. 全连接层：分类

输入：EEG 数据 shape=(channels, time_steps, freq_bins)
输出：风险等级分类
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class Conv3DBlock(nn.Module):
    """3D 卷积块：Conv3D + BN + ReLU + Dropout"""
    
    def __init__(self, in_channels, out_channels,
                 kernel_size=(3, 3, 3), 
                 stride=(1, 1, 1),
                 padding=(1, 1, 1),
                 dropout=0.3):
        super().__init__()
        self.conv = nn.Conv3d(
            in_channels, out_channels,
            kernel_size, stride, padding
        )
        self.bn = nn.BatchNorm3d(out_channels)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
        self.pool = nn.MaxPool3d((1, 2, 2))
        
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.pool(x)
        return x


class CRNNModel(nn.Module):
    """
    3D-CRNN 模型
    
    论文 Section 3.3 - PEDS 策略核心
    
    输入 EEG 格式：
    - channels: 电极通道数（如 32）
    - time_steps: 时间步数
    - freq_bins: 频率bins（如 STFT 后的 64）
    
    架构：
    3D Conv blocks → Reshape → BiLSTM → FC → Output
    """
    
    def __init__(self, num_channels=32, num_classes=4,
                 num_filters=16, hidden_size=128,
                 num_layers=2, dropout=0.3):
        super().__init__()
        
        # 3D 卷积特征提取
        self.conv_blocks = nn.ModuleList([
            Conv3DBlock(1, num_filters,
                       dropout=dropout),
            Conv3DBlock(num_filters, num_filters*2,
                       dropout=dropout),
            Conv3DBlock(num_filters*2, num_filters*4,
                       dropout=dropout),
        ])
        
        # 计算 LSTM 输入维度
        self.lstm_input_size = num_filters * 4
        
        # 双向 LSTM
        self.lstm = nn.LSTM(
            input_size=self.lstm_input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,
            dropout=dropout if num_layers > 1 else 0
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(hidden_size * 2, 64),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(64, num_classes)
        )
        
    def forward(self, x):
        """
        前向传播
        
        Args:
            x: EEG 数据, shape=(B, 1, C, T, F)
               B=batch, C=channels, T=time, F=freq
            
        Returns:
            logits: shape=(B, num_classes)
        """
        # 3D 卷积
        for block in self.conv_blocks:
            x = block(x)
        
        # x shape: (B, num_filters*4, C', T', F')
        # Reshape 为 (B, T', num_filters*4*C'*F')
        B, C_out, T_out, F_out = x.shape[0], x.shape[1], \
                                 x.shape[2], x.shape[3]
        
        # 通道和频率维合并，时间维作为序列
        x = x.permute(0, 2, 1, 3)  # (B, T, C, F)
        x = x.reshape(B, T_out, -1)  # (B, T, C*F)
        
        # BiLSTM
        lstm_out, _ = self.lstm(x)
        
        # 取最后时间步的输出
        out = lstm_out[:, -1, :]
        
        # 分类
        logits = self.classifier(out)
        return logits


# ========== 测试 ==========

if __name__ == "__main__":
    # 模型参数
    model = CRNNModel(
        num_channels=32,
        num_classes=4,
        num_filters=16,
        hidden_size=128,
        num_layers=2,
        dropout=0.3
    )
    
    # 模拟输入：batch=4, channel=1, EEG=32ch, 100时间步, 64频率bin
    x = torch.randn(4, 1, 32, 100, 64)
    
    # 前向传播
    logits = model(x)
    
    print(f"=== 3D-CRNN 模型 ===")
    print(f"  输入 shape: {x.shape}")
    print(f"  输出 shape: {logits.shape}")
    print(f"  参数量: {sum(p.numel() for p in model.parameters()):,}")
    
    # 预测
    predictions = F.softmax(logits, dim=-1)
    predicted_classes = torch.argmax(predictions, dim=-1)
    print(f"  预测类别: {predicted_classes.tolist()}")
    print(f"  置信度: {predictions.max(dim=-1).values.tolist()}")