#长短时记忆网络 (LSTM/GRU)：解决梯度消失，捕捉长距离依赖

import torch
import torch.nn as nn

class BiLSTMTextClassifier(nn.Module):
    """
    双向LSTM文本分类器
    适用于情感分析、新闻分类等短文本任务
    """
    def __init__(
        self,
        vocab_size: int,    # 词表大小
        embed_dim: int = 256,  # 词嵌入维度
        hidden_dim: int = 256, # LSTM隐藏层维度
        num_layers: int = 2,   # LSTM层数
        dropout: float = 0.3,  # dropout比例，防止过拟合
        num_classes: int = 2    # 分类类别数（二分类：积极/消极）
    ):
        super().__init__()
        
        # 1. 词嵌入层：把词ID转换成低维稠密向量
        self.embedding = nn.Embedding(
            vocab_size, embed_dim, padding_idx=0  # padding_idx=0：忽略词表中的填充词
        )
        
        # 2. 双向LSTM层
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,   # 输入输出的第一维度是batch_size（更符合习惯）
            bidirectional=True, # 双向：前向看前文，后向看后文
            dropout=dropout if num_layers > 1 else 0  # 只有多层LSTM才加层间dropout
        )
        
        # 3. 全连接分类头：双向拼接后的维度是 hidden_dim*2
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_classes)
        )

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        前向传播
        input_ids: (batch_size, seq_len) → 每个样本是词ID序列
        logits: (batch_size, num_classes) → 每个样本对应类别的未归一化分数
        """
        # 词嵌入：(B, L) → (B, L, E)
        embedded = self.embedding(input_ids)
        
        # LSTM计算
        # output: (B, L, 2H) → 每个时刻的双向隐藏状态拼接
        # (h_n, c_n): 最后时刻的隐藏状态和细胞状态，(2*num_layers, B, H)
        output, (h_n, _) = self.lstm(embedded)
        
        # 取最后一层的双向隐藏状态拼接 → (B, 2H)
        last_layer_forward = h_n[-2]  # 前向最后一层的最后一个隐藏状态
        last_layer_backward = h_n[-1] # 后向最后一层的最后一个隐藏状态
        final_hidden = torch.cat([last_layer_forward, last_layer_backward], dim=-1)
        
        # 分类头计算
        logits = self.classifier(final_hidden)
        return logits

import torch
import torch.nn as nn

class BiGRUTextClassifier(nn.Module):
    """
    双向GRU文本分类器
    轻量高效，适合快速原型验证或简单任务
    """
    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 256,
        hidden_dim: int = 256,
        num_layers: int = 2,
        dropout: float = 0.3,
        num_classes: int = 2
    ):
        super().__init__()
        
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        
        # 替换成nn.GRU
        self.gru = nn.GRU(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,
            dropout=dropout if num_layers > 1 else 0
        )
        
        # 分类头不变
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_classes)
        )

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        embedded = self.embedding(input_ids)
        # GRU只返回output和h_n，没有c_n
        _, h_n = self.gru(embedded)
        final_hidden = torch.cat([h_n[-2], h_n[-1]], dim=-1)
        logits = self.classifier(final_hidden)
        return logits

import torch
from torch.utils.data import DataLoader

def train_one_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
    clip_max_norm: float = 1.0  # 梯度裁剪的最大范数
) -> tuple[float, float]:
    """
    训练单轮模型
    返回：平均损失、平均准确率
    """
    model.train()
    total_loss = 0.0
    correct_preds = 0
    total_preds = 0

    for batch in dataloader:
        # 把数据移到GPU/CPU
        input_ids = batch["input_ids"].to(device)
        labels = batch["label"].to(device)

        # 前向传播
        optimizer.zero_grad()
        logits = model(input_ids)
        loss = criterion(logits, labels)

        # 反向传播 + 梯度裁剪 + 参数更新
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip_max_norm)
        optimizer.step()

        # 统计指标
        total_loss += loss.item()
        preds = logits.argmax(dim=-1)
        correct_preds += (preds == labels).sum().item()
        total_preds += labels.size(0)

    avg_loss = total_loss / len(dataloader)
    avg_acc = correct_preds / total_preds
    return avg_loss, avg_acc

import torch

def predict_sentiment(
    model: nn.Module,
    tokenizer,  # 假设已初始化好的Tokenizer
    text: str,
    device: torch.device
) -> dict[str, float]:
    """
    预测单条文本的情感
    返回：积极、消极的概率字典
    """
    model.eval()
    with torch.no_grad():  # 推理时不需要计算梯度
        # 分词+转ID+补填充（这里简化，实际用tokenizer的__call__更方便）
        tokens = tokenizer.tokenize(text)
        ids = tokenizer.convert_tokens_to_ids(tokens)
        # 加batch维度 → (1, seq_len)
        input_ids = torch.tensor([ids]).to(device)
        
        # 前向传播
        logits = model(input_ids)
        # 用softmax把未归一化分数转成概率
        probs = torch.softmax(logits, dim=-1).squeeze(0)  # 去掉batch维度
    
    return {
        "positive": round(probs[1].item(), 4),
        "negative": round(probs[0].item(), 4)
    }