import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from config import Config
from model import make_model
from dataset import TranslationDataset, collate_fn

def train_model():
    """模型训练主函数"""
    config = Config()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 1. 加载数据集
    train_dataset = TranslationDataset(config.train_data)
    valid_dataset = TranslationDataset(config.valid_data)
    
    train_loader = DataLoader(
        train_dataset, batch_size=config.batch_size, 
        shuffle=True, collate_fn=collate_fn
    )
    valid_loader = DataLoader(
        valid_dataset, batch_size=config.batch_size, 
        shuffle=False, collate_fn=collate_fn
    )
    
    # 2. 创建模型
    model = make_model(
        src_vocab=config.src_vocab_size,
        tgt_vocab=config.tgt_vocab_size,
        N=config.n_layers,
        d_model=config.d_model,
        d_ff=config.d_ff,
        h=config.n_heads,
        dropout=config.dropout
    ).to(device)
    
    # 3. 损失函数 (忽略padding位置)
    criterion = nn.CrossEntropyLoss(
        ignore_index=config.pad_idx, 
        label_smoothing=0.1
    )
    
    # 4. 优化器
    optimizer = torch.optim.Adam(
        model.parameters(), 
        lr=config.learning_rate,
        betas=(0.9, 0.98),
        eps=1e-9
    )
    
    # 5. 学习率调度器 (可选: warmup)
    scheduler = torch.optim.lr_scheduler.LambdaLR(
        optimizer, lr_lambda=lambda step: warmup_lr(step, config)
    )
    
    # 6. 训练循环
    best_bleu = 0
    for epoch in range(config.epochs):
        # 训练阶段
        model.train()
        total_loss = 0
        for batch in train_loader:
            src = batch.src.to(device)
            tgt = batch.tgt.to(device)
            src_mask = batch.src_mask.to(device)
            tgt_mask = batch.tgt_mask.to(device)
            
            # 前向传播
            output = model(src, tgt[:, :-1], src_mask, tgt_mask)
            
            # 计算损失
            loss = criterion(
                output.reshape(-1, output.size(-1)),
                tgt[:, 1:].reshape(-1)
            )
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()
            
            total_loss += loss.item()
        
        avg_loss = total_loss / len(train_loader)
        
        # 验证阶段
        val_loss, bleu_score = evaluate(model, valid_loader, criterion, device)
        
        print(f"Epoch {epoch+1}/{config.epochs}")
        print(f"  Train Loss: {avg_loss:.4f}")
        print(f"  Val Loss: {val_loss:.4f}")
        print(f"  BLEU Score: {bleu_score:.2f}")
        
        # 保存最佳模型
        if bleu_score > best_bleu:
            best_bleu = bleu_score
            torch.save(model.state_dict(), 'runs/best.pt')
            print(f"  → 保存最佳模型 (BLEU: {bleu_score:.2f})")
    
    # 保存最终模型
    torch.save(model.state_dict(), 'runs/last.pt')
    print(f"训练完成! 最佳BLEU: {best_bleu:.2f}")


def warmup_lr(step, config):
    """Warmup学习率调度"""
    warmup_steps = 4000
    if step < warmup_steps:
        return step / warmup_steps
    return 1.0 / (step ** 0.5)


if __name__ == "__main__":
    train_model()