#!/usr/bin/env python3
"""
Mini-BERT SFT微调脚本
基于预训练的Mini-BERT模型在SST-2任务上进行监督微调
"""

import os
import json
import argparse
import time
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torch.cuda.amp import autocast, GradScaler
from transformers import (
    BertConfig, 
    BertForSequenceClassification, 
    BertTokenizer,
    AdamW,
    get_linear_schedule_with_warmup
)
import pandas as pd
from sklearn.metrics import accuracy_score, f1_score, classification_report, confusion_matrix
from tqdm import tqdm
import numpy as np
from collections import deque

class SST2Dataset(Dataset):
    """SST-2数据集类"""
    
    def __init__(self, texts, labels, tokenizer, max_length=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        text = str(self.texts[idx])
        label = self.labels[idx]
        
        # 分词和编码
        encoding = self.tokenizer(
            text,
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors='pt'
        )
        
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'token_type_ids': encoding.get('token_type_ids', torch.zeros_like(encoding['input_ids'])).flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }

class SFTTrainingMonitor:
    """SFT训练监控器"""
    
    def __init__(self, patience=3):
        self.patience = patience
        self.best_val_acc = 0
        self.epochs_without_improvement = 0
        self.training_history = []
        
    def update(self, epoch, train_loss, train_acc, val_loss, val_acc, val_f1):
        """更新训练记录"""
        self.training_history.append({
            'epoch': epoch,
            'train_loss': train_loss,
            'train_acc': train_acc,
            'val_loss': val_loss,
            'val_acc': val_acc,
            'val_f1': val_f1
        })
        
        # 检查是否有改善
        if val_acc > self.best_val_acc:
            self.best_val_acc = val_acc
            self.epochs_without_improvement = 0
            return True  # 有改善
        else:
            self.epochs_without_improvement += 1
            return False  # 无改善
    
    def should_early_stop(self):
        """判断是否应该早停"""
        return self.epochs_without_improvement >= self.patience
    
    def print_epoch_summary(self, epoch, train_loss, train_acc, val_loss, val_acc, val_f1, lr):
        """打印epoch总结"""
        print(f"\n📊 Epoch {epoch} 总结:")
        print(f"   训练损失: {train_loss:.4f} | 训练准确率: {train_acc:.4f} ({train_acc*100:.2f}%)")
        print(f"   验证损失: {val_loss:.4f} | 验证准确率: {val_acc:.4f} ({val_acc*100:.2f}%)")
        print(f"   验证F1: {val_f1:.4f} | 学习率: {lr:.2e}")
        print(f"   最佳验证准确率: {self.best_val_acc:.4f} | 无改善轮数: {self.epochs_without_improvement}")

def init_distributed():
    """初始化分布式训练"""
    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
        torch.distributed.init_process_group(backend="nccl")
        local_rank = int(os.environ.get("LOCAL_RANK", "0"))
        torch.cuda.set_device(local_rank)
        return True, local_rank
    return False, 0

def load_pretrained_model(pretrained_path, tokenizer_path, num_labels=2):
    """加载预训练模型"""
    print(f"🔄 加载预训练模型: {pretrained_path}")
    
    # 加载分词器
    try:
        tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
        print(f"✅ 成功加载分词器")
    except Exception as e:
        print(f"⚠️  加载分词器失败: {e}")
        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
        print(f"🔄 使用默认分词器")
    
    # 加载配置
    config_path = os.path.join(pretrained_path, "config.json")
    if os.path.exists(config_path):
        config = BertConfig.from_pretrained(pretrained_path)
    else:
        # 使用默认配置
        config = BertConfig(
            vocab_size=tokenizer.vocab_size,
            hidden_size=512,
            num_hidden_layers=6,
            num_attention_heads=8,
            intermediate_size=2048,
            max_position_embeddings=512,
            type_vocab_size=2,
            hidden_dropout_prob=0.1,
            attention_probs_dropout_prob=0.1,
            layer_norm_eps=1e-12,
            hidden_act="gelu"
        )
    
    # 设置分类标签数
    config.num_labels = num_labels
    
    # 创建分类模型
    model = BertForSequenceClassification(config)
    
    # 加载预训练权重
    try:
        # 查找权重文件
        weight_files = ["pytorch_model.bin", "model.safetensors", "pytorch_model.safetensors"]
        weight_path = None
        
        for weight_file in weight_files:
            candidate_path = os.path.join(pretrained_path, weight_file)
            if os.path.exists(candidate_path):
                weight_path = candidate_path
                break
        
        if weight_path is None:
            raise FileNotFoundError("未找到权重文件")
        
        # 加载权重
        if weight_path.endswith('.safetensors'):
            from safetensors.torch import load_file
            pretrained_dict = load_file(weight_path)
        else:
            pretrained_dict = torch.load(weight_path, map_location="cpu")
        
        # 处理分布式训练的权重（移除module.前缀）
        if any(key.startswith('module.') for key in pretrained_dict.keys()):
            pretrained_dict = {key.replace('module.', ''): value for key, value in pretrained_dict.items()}
        
        # 获取当前模型的状态字典
        model_dict = model.state_dict()
        
        # 过滤出匹配的权重（排除分类头）
        filtered_dict = {}
        for k, v in pretrained_dict.items():
            if k in model_dict and "classifier" not in k:
                filtered_dict[k] = v
        
        # 更新模型权重
        model_dict.update(filtered_dict)
        model.load_state_dict(model_dict)
        
        print(f"✅ 成功加载 {len(filtered_dict)} 个预训练参数")
        print(f"🆕 新初始化分类头参数")
        
    except Exception as e:
        print(f"⚠️  加载预训练权重失败: {e}")
        print(f"🔄 使用随机初始化的模型")
    
    return model, tokenizer, config

def load_sst2_data(train_path, val_path, test_path=None):
    """加载SST-2数据"""
    print(f"📊 加载SST-2数据...")
    
    # 加载训练数据
    if train_path.endswith('.csv'):
        train_df = pd.read_csv(train_path)
    elif train_path.endswith('.jsonl'):
        train_df = pd.read_json(train_path, lines=True)
    else:
        raise ValueError("不支持的文件格式")
    
    # 加载验证数据
    if val_path.endswith('.csv'):
        val_df = pd.read_csv(val_path)
    elif val_path.endswith('.jsonl'):
        val_df = pd.read_json(val_path, lines=True)
    else:
        raise ValueError("不支持的文件格式")
    
    # 加载测试数据（可选）
    test_df = None
    if test_path:
        if test_path.endswith('.csv'):
            test_df = pd.read_csv(test_path)
        elif test_path.endswith('.jsonl'):
            test_df = pd.read_json(test_path, lines=True)
    
    print(f"✅ 数据加载完成:")
    print(f"   训练集: {len(train_df)} 样本")
    print(f"   验证集: {len(val_df)} 样本")
    if test_df is not None:
        print(f"   测试集: {len(test_df)} 样本")
    
    return train_df, val_df, test_df

def create_data_loaders(train_df, val_df, tokenizer, batch_size=16, max_length=128, test_df=None):
    """创建数据加载器"""
    # 确定文本和标签列名
    text_col = 'sentence' if 'sentence' in train_df.columns else 'text'
    label_col = 'label'
    
    # 创建数据集
    train_dataset = SST2Dataset(
        train_df[text_col].tolist(),
        train_df[label_col].tolist(),
        tokenizer,
        max_length
    )
    
    val_dataset = SST2Dataset(
        val_df[text_col].tolist(),
        val_df[label_col].tolist(),
        tokenizer,
        max_length
    )
    
    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size * 2,  # 验证时可以用更大的batch
        shuffle=False,
        num_workers=4,
        pin_memory=True
    )
    
    test_loader = None
    if test_df is not None:
        test_dataset = SST2Dataset(
            test_df[text_col].tolist(),
            test_df[label_col].tolist(),
            tokenizer,
            max_length
        )
        test_loader = DataLoader(
            test_dataset,
            batch_size=batch_size * 2,
            shuffle=False,
            num_workers=4,
            pin_memory=True
        )
    
    return train_loader, val_loader, test_loader

def setup_optimizer_and_scheduler(model, train_loader, learning_rate=2e-5, weight_decay=0.01, num_epochs=5, warmup_ratio=0.1):
    """设置优化器和调度器"""
    # 分组权重衰减
    no_decay = ['bias', 'LayerNorm.weight', 'LayerNorm.bias']
    optimizer_grouped_parameters = [
        {
            'params': [p for n, p in model.named_parameters() 
                      if not any(nd in n for nd in no_decay)],
            'weight_decay': weight_decay
        },
        {
            'params': [p for n, p in model.named_parameters() 
                      if any(nd in n for nd in no_decay)],
            'weight_decay': 0.0
        }
    ]
    
    optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate, eps=1e-8)
    
    # 学习率调度器
    total_steps = len(train_loader) * num_epochs
    warmup_steps = int(total_steps * warmup_ratio)
    
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=total_steps
    )
    
    return optimizer, scheduler

def train_epoch(model, train_loader, optimizer, scheduler, scaler, device):
    """训练一个epoch"""
    model.train()
    total_loss = 0
    predictions = []
    true_labels = []
    
    progress_bar = tqdm(train_loader, desc="训练")
    
    for batch in progress_bar:
        # 将数据移到GPU
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        token_type_ids = batch['token_type_ids'].to(device)
        labels = batch['labels'].to(device)
        
        # 前向传播
        with autocast():
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                token_type_ids=token_type_ids,
                labels=labels
            )
            loss = outputs.loss
            logits = outputs.logits
        
        # 反向传播
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()
        
        # 记录结果
        total_loss += loss.item()
        batch_predictions = torch.argmax(logits, dim=-1)
        predictions.extend(batch_predictions.cpu().numpy())
        true_labels.extend(labels.cpu().numpy())
        
        # 更新进度条
        progress_bar.set_postfix({
            'loss': f'{loss.item():.4f}',
            'lr': f'{scheduler.get_last_lr()[0]:.2e}'
        })
    
    avg_loss = total_loss / len(train_loader)
    accuracy = accuracy_score(true_labels, predictions)
    
    return avg_loss, accuracy

def evaluate_model(model, dataloader, device):
    """评估模型"""
    model.eval()
    total_loss = 0
    predictions = []
    true_labels = []
    
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="评估"):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            token_type_ids = batch['token_type_ids'].to(device)
            labels = batch['labels'].to(device)
            
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                token_type_ids=token_type_ids,
                labels=labels
            )
            
            loss = outputs.loss
            logits = outputs.logits
            
            total_loss += loss.item()
            batch_predictions = torch.argmax(logits, dim=-1)
            predictions.extend(batch_predictions.cpu().numpy())
            true_labels.extend(labels.cpu().numpy())
    
    avg_loss = total_loss / len(dataloader)
    accuracy = accuracy_score(true_labels, predictions)
    f1 = f1_score(true_labels, predictions)
    
    return avg_loss, accuracy, f1, predictions, true_labels

def save_model(model, tokenizer, output_dir, epoch, is_best=False):
    """保存模型"""
    if is_best:
        save_dir = os.path.join(output_dir, "best_model")
    else:
        save_dir = os.path.join(output_dir, f"epoch_{epoch}")
    
    os.makedirs(save_dir, exist_ok=True)
    
    # 保存模型
    if hasattr(model, 'module'):
        model.module.save_pretrained(save_dir)
    else:
        model.save_pretrained(save_dir)
    
    # 保存分词器
    tokenizer.save_pretrained(save_dir)
    
    print(f"💾 模型已保存到: {save_dir}")

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="Mini-BERT SFT微调")
    
    # 模型和数据路径
    parser.add_argument("--pretrained_model_path", required=True, help="预训练模型路径")
    parser.add_argument("--tokenizer_path", required=True, help="分词器路径")
    parser.add_argument("--train_data", required=True, help="训练数据路径")
    parser.add_argument("--val_data", required=True, help="验证数据路径")
    parser.add_argument("--test_data", help="测试数据路径（可选）")
    parser.add_argument("--output_dir", required=True, help="输出目录")
    
    # 训练参数
    parser.add_argument("--num_epochs", type=int, default=5, help="训练轮数")
    parser.add_argument("--batch_size", type=int, default=16, help="批大小")
    parser.add_argument("--learning_rate", type=float, default=2e-5, help="学习率")
    parser.add_argument("--weight_decay", type=float, default=0.01, help="权重衰减")
    parser.add_argument("--warmup_ratio", type=float, default=0.1, help="预热比例")
    parser.add_argument("--max_length", type=int, default=128, help="最大序列长度")
    parser.add_argument("--patience", type=int, default=3, help="早停耐心值")
    
    # 其他参数
    parser.add_argument("--seed", type=int, default=42, help="随机种子")
    parser.add_argument("--fp16", action="store_true", help="使用混合精度训练")
    
    return parser.parse_args()

def main():
    args = parse_args()
    
    # 设置随机种子
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    
    # 初始化分布式训练
    is_dist, local_rank = init_distributed()
    device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
    is_main = (not is_dist) or (int(os.environ.get("RANK", "0")) == 0)
    
    if is_main:
        print("🚀 开始Mini-BERT SFT微调")
        print(f"📊 训练参数:")
        print(f"   - 预训练模型: {args.pretrained_model_path}")
        print(f"   - 训练数据: {args.train_data}")
        print(f"   - 验证数据: {args.val_data}")
        print(f"   - 输出目录: {args.output_dir}")
        print(f"   - 训练轮数: {args.num_epochs}")
        print(f"   - 批大小: {args.batch_size}")
        print(f"   - 学习率: {args.learning_rate}")
        print(f"   - 设备: {device}")
        print("=" * 60)
    
    # 加载模型和分词器
    model, tokenizer, config = load_pretrained_model(
        args.pretrained_model_path, 
        args.tokenizer_path, 
        num_labels=2
    )
    model.to(device)
    
    # 分布式训练包装
    if is_dist:
        model = torch.nn.parallel.DistributedDataParallel(
            model, device_ids=[local_rank], output_device=local_rank
        )
    
    # 加载数据
    train_df, val_df, test_df = load_sst2_data(args.train_data, args.val_data, args.test_data)
    
    # 创建数据加载器
    train_loader, val_loader, test_loader = create_data_loaders(
        train_df, val_df, tokenizer, args.batch_size, args.max_length, test_df
    )
    
    # 设置优化器和调度器
    optimizer, scheduler = setup_optimizer_and_scheduler(
        model, train_loader, args.learning_rate, args.weight_decay, 
        args.num_epochs, args.warmup_ratio
    )
    
    # 混合精度训练
    scaler = GradScaler(enabled=args.fp16)
    
    # 训练监控器
    monitor = SFTTrainingMonitor(patience=args.patience)
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 开始训练
    if is_main:
        print("\n🎯 开始训练循环")
    
    for epoch in range(args.num_epochs):
        if is_main:
            print(f"\n📅 Epoch {epoch + 1}/{args.num_epochs}")
        
        # 训练
        train_loss, train_acc = train_epoch(model, train_loader, optimizer, scheduler, scaler, device)
        
        # 验证
        val_loss, val_acc, val_f1, val_preds, val_labels = evaluate_model(model, val_loader, device)
        
        # 更新监控器
        is_best = monitor.update(epoch + 1, train_loss, train_acc, val_loss, val_acc, val_f1)
        
        if is_main:
            # 打印epoch总结
            monitor.print_epoch_summary(
                epoch + 1, train_loss, train_acc, val_loss, val_acc, val_f1, 
                scheduler.get_last_lr()[0]
            )
            
            # 保存模型
            save_model(model, tokenizer, args.output_dir, epoch + 1, is_best)
            
            # 检查早停
            if monitor.should_early_stop():
                print(f"\n⏰ 早停触发！连续 {args.patience} 轮无改善")
                break
    
    # 最终评估
    if is_main:
        print(f"\n🎉 训练完成！")
        print(f"📊 最终结果:")
        print(f"   - 最佳验证准确率: {monitor.best_val_acc:.4f} ({monitor.best_val_acc*100:.2f}%)")
        
        # 如果有测试集，进行测试
        if test_loader is not None:
            print(f"\n🧪 在测试集上评估...")
            test_loss, test_acc, test_f1, test_preds, test_labels = evaluate_model(model, test_loader, device)
            print(f"📊 测试集结果:")
            print(f"   - 测试准确率: {test_acc:.4f} ({test_acc*100:.2f}%)")
            print(f"   - 测试F1: {test_f1:.4f}")
            
            # 详细分类报告
            print(f"\n📋 详细分类报告:")
            print(classification_report(test_labels, test_preds, target_names=["负面", "正面"]))
        
        # 保存训练历史
        history_path = os.path.join(args.output_dir, "training_history.json")
        with open(history_path, 'w') as f:
            json.dump(monitor.training_history, f, indent=2)
        print(f"📈 训练历史已保存到: {history_path}")
    
    # 清理分布式训练
    if is_dist:
        torch.distributed.destroy_process_group()

if __name__ == "__main__":
    main()