#!/usr/bin/env python3
"""
BERT情感分类任务评估脚本
基于SFT章节的eval_qwen.py，适配BERT模型进行情感分析评估
"""

import os
import json
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from transformers import (
    BertForSequenceClassification, 
    BertTokenizer, 
    BertConfig,
    AutoTokenizer,
    AutoModelForSequenceClassification
)
import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Optional
import argparse
from tqdm import tqdm
import psutil
import GPUtil
from datetime import datetime
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, classification_report
import glob
import re

class BERTSentimentEvaluator:
    """BERT情感分类评估器"""
    
    def __init__(self, model_path: str, tokenizer_path: str = None, device: str = "cuda"):
        self.device = device if torch.cuda.is_available() else "cpu"
        self.model_path = model_path
        self.tokenizer_path = tokenizer_path or model_path
        
        # 加载模型和分词器
        self.model, self.tokenizer = self._load_model_and_tokenizer()
        self.model.eval()
        
        # 评估指标存储
        self.metrics = {
            'accuracy': 0.0,
            'f1_score': 0.0,
            'precision': 0.0,
            'recall': 0.0,
            'confusion_matrix': None,
            'inference_times': [],
            'predictions': [],
            'true_labels': []
        }
        
    def _find_latest_checkpoint(self, model_path: str) -> str:
        """找到最新的检查点"""
        print(f"🔍 查找最新检查点...")
        
        # 检查是否是直接的模型目录
        if os.path.exists(os.path.join(model_path, "config.json")):
            return model_path
        
        # 查找所有step-*文件夹
        step_dirs = []
        for item in os.listdir(model_path):
            if item.startswith('step-') and os.path.isdir(os.path.join(model_path, item)):
                try:
                    step_num = int(item.split('-')[1])
                    step_dirs.append((step_num, item))
                except ValueError:
                    continue
        
        if not step_dirs:
            raise FileNotFoundError(f"未找到任何检查点文件夹在: {model_path}")
        
        # 按步数排序，取最新的
        step_dirs.sort(key=lambda x: x[0], reverse=True)
        latest_step, latest_dir = step_dirs[0]
        
        latest_path = os.path.join(model_path, latest_dir)
        print(f"✅ 找到最新检查点: {latest_dir} (步数: {latest_step})")
        
        return latest_path
    
    def _load_model_and_tokenizer(self):
        """加载BERT分类模型和分词器"""
        print(f"🔄 加载模型从: {self.model_path}")
        print(f"🔄 加载分词器从: {self.tokenizer_path}")
        
        # 检查模型路径是否存在
        if not os.path.exists(self.model_path):
            raise FileNotFoundError(f"模型路径不存在: {self.model_path}")
        
        if not os.path.exists(self.tokenizer_path):
            raise FileNotFoundError(f"分词器路径不存在: {self.tokenizer_path}")
        
        # 找到最新检查点
        checkpoint_path = self._find_latest_checkpoint(self.model_path)
        
        # 加载分词器
        try:
            tokenizer = BertTokenizer.from_pretrained(self.tokenizer_path)
        except:
            # 如果BERT分词器失败，尝试AutoTokenizer
            tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path)
        
        # 检查是否有config.json
        config_path = os.path.join(checkpoint_path, "config.json")
        if os.path.exists(config_path):
            try:
                config = BertConfig.from_pretrained(checkpoint_path)
                print("✅ 从检查点加载BERT配置")
            except:
                # 如果BERT配置失败，尝试AutoConfig
                from transformers import AutoConfig
                config = AutoConfig.from_pretrained(checkpoint_path)
                print("✅ 从检查点加载AutoConfig")
        else:
            # 如果没有config.json，使用默认配置
            print("⚠️  未找到config.json，使用默认配置")
            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,
                num_labels=2  # 二分类任务
            )
        
        # 加载预训练模型
        try:
            # 首先尝试BertForSequenceClassification
            model = BertForSequenceClassification.from_pretrained(checkpoint_path, config=config)
            print("✅ 从检查点加载BERT分类模型成功")
        except Exception as e:
            print(f"⚠️  BERT分类模型加载失败: {e}")
            print("🔄 尝试AutoModelForSequenceClassification...")
            try:
                model = AutoModelForSequenceClassification.from_pretrained(checkpoint_path, config=config)
                print("✅ 从检查点加载Auto分类模型成功")
            except Exception as e2:
                print(f"⚠️  Auto分类模型加载失败: {e2}")
                print("🔄 尝试创建新模型并加载权重...")
                
                # 创建新模型
                model = BertForSequenceClassification(config)
                
                # 尝试加载权重
                weight_loaded = False
                checkpoint_files = os.listdir(checkpoint_path)
                model_files = [f for f in checkpoint_files if f.endswith(('.bin', '.safetensors', '.pt', '.pth'))]
                
                for model_file in model_files:
                    try:
                        weight_path = os.path.join(checkpoint_path, model_file)
                        if model_file.endswith('.safetensors'):
                            from safetensors.torch import load_file
                            state_dict = load_file(weight_path)
                        else:
                            state_dict = torch.load(weight_path, map_location='cpu')
                        
                        # 处理分布式训练的权重（移除module.前缀）
                        if any(key.startswith('module.') for key in state_dict.keys()):
                            state_dict = {key.replace('module.', ''): value for key, value in state_dict.items()}
                        
                        model.load_state_dict(state_dict, strict=False)
                        weight_loaded = True
                        print(f"✅ 成功加载权重: {model_file}")
                        break
                    except Exception as e3:
                        print(f"⚠️  加载权重失败 {model_file}: {e3}")
                        continue
                
                if not weight_loaded:
                    raise RuntimeError("无法加载任何模型权重文件")
        
        model.to(self.device)
        
        print(f"✅ 模型加载完成")
        print(f"   - 词汇表大小: {config.vocab_size}")
        print(f"   - 隐藏层维度: {config.hidden_size}")
        print(f"   - 层数: {config.num_hidden_layers}")
        print(f"   - 注意力头数: {config.num_attention_heads}")
        print(f"   - 分类标签数: {config.num_labels}")
        
        return model, tokenizer
    
    def load_evaluation_data(self, data_path: str) -> Tuple[List[str], List[int]]:
        """加载评估数据"""
        print(f"📊 加载评估数据: {data_path}")
        
        if data_path.endswith('.csv'):
            # CSV格式数据
            df = pd.read_csv(data_path)
            if 'text' in df.columns and 'label' in df.columns:
                texts = df['text'].tolist()
                labels = df['label'].tolist()
            else:
                raise ValueError("CSV文件必须包含'text'和'label'列")
        elif data_path.endswith('.jsonl'):
            # JSONL格式数据
            texts = []
            labels = []
            with open(data_path, 'r', encoding='utf-8') as f:
                for line in f:
                    data = json.loads(line.strip())
                    if 'text' in data and 'label' in data:
                        texts.append(data['text'])
                        labels.append(data['label'])
                    elif 'sentence' in data and 'label' in data:
                        texts.append(data['sentence'])
                        labels.append(data['label'])
                    else:
                        print(f"⚠️  跳过无效数据行: {data}")
            print(f"✅ 从JSONL加载了 {len(texts)} 个样本")
        else:
            raise ValueError("不支持的数据格式，请使用CSV或JSONL格式")
        
        print(f"📊 数据统计:")
        print(f"   - 总样本数: {len(texts)}")
        print(f"   - 标签分布: {dict(pd.Series(labels).value_counts().sort_index())}")
        
        return texts, labels
    
    def predict_sentiment(self, text: str) -> Tuple[int, float]:
        """预测单条文本的情感"""
        # 编码输入
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=512
        ).to(self.device)
        
        # 前向传播
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = F.softmax(logits, dim=-1)
            predicted_class = torch.argmax(logits, dim=-1).item()
            confidence = probabilities[0][predicted_class].item()
        
        return predicted_class, confidence
    
    def evaluate_model(self, texts: List[str], true_labels: List[int]) -> Dict[str, any]:
        """评估模型性能"""
        print(f"\n🔍 开始评估...")
        print("=" * 60)
        
        predictions = []
        confidences = []
        inference_times = []
        
        # 逐条预测
        for i, text in enumerate(tqdm(texts, desc="情感分析评估")):
            start_time = time.time()
            
            try:
                pred, confidence = self.predict_sentiment(text)
                predictions.append(pred)
                confidences.append(confidence)
                
                inference_time = time.time() - start_time
                inference_times.append(inference_time)
                
            except Exception as e:
                print(f"⚠️  样本 {i+1} 预测失败: {e}")
                predictions.append(0)  # 默认预测负面
                confidences.append(0.0)
                inference_times.append(0)
        
        # 计算指标
        accuracy = accuracy_score(true_labels, predictions)
        f1 = f1_score(true_labels, predictions)
        cm = confusion_matrix(true_labels, predictions)
        
        # 详细分类报告
        report = classification_report(
            true_labels, predictions,
            target_names=["负面", "正面"],
            output_dict=True
        )
        
        # 存储结果
        self.metrics.update({
            'accuracy': accuracy,
            'f1_score': f1,
            'precision': report['正面']['precision'],
            'recall': report['正面']['recall'],
            'confusion_matrix': cm,
            'inference_times': inference_times,
            'predictions': predictions,
            'true_labels': true_labels,
            'confidences': confidences
        })
        
        return self.metrics
    
    def print_evaluation_results(self):
        """打印评估结果"""
        metrics = self.metrics
        cm = metrics['confusion_matrix']
        
        print("\n" + "=" * 60)
        print("BERT情感分类模型评估报告")
        print("=" * 60)
        print(f"模型路径: {self.model_path}")
        print(f"分词器路径: {self.tokenizer_path}")
        print(f"评测样本数: {len(metrics['true_labels'])}")
        print(f"任务: 电影评论情感分析 (二分类)")
        print()
        
        # 性能指标
        print("📊 性能指标:")
        print(f"  准确率 (Accuracy): {metrics['accuracy']:.3f} ({metrics['accuracy']*100:.1f}%)")
        print(f"  F1分数: {metrics['f1_score']:.3f}")
        print(f"  精确率 (正面): {metrics['precision']:.3f}")
        print(f"  召回率 (正面): {metrics['recall']:.3f}")
        print()
        
        # 混淆矩阵
        print("🔍 混淆矩阵:")
        print("        预测")
        print("实际    负面  正面")
        print(f"负面    {cm[0,0]:4d}  {cm[0,1]:4d}")
        print(f"正面    {cm[1,0]:4d}  {cm[1,1]:4d}")
        print()
        
        # 推理性能
        avg_time = np.mean(metrics['inference_times'])
        total_time = sum(metrics['inference_times'])
        throughput = len(metrics['predictions']) / total_time if total_time > 0 else 0
        
        print("⚡ 推理性能:")
        print(f"  平均推理时间: {avg_time:.4f}秒/样本")
        print(f"  总推理时间: {total_time:.1f}秒")
        print(f"  吞吐量: {throughput:.2f}样本/秒")
        print()
        
        # 置信度分析
        avg_confidence = np.mean(metrics['confidences'])
        high_confidence = sum(1 for c in metrics['confidences'] if c > 0.8)
        
        print("🎯 置信度分析:")
        print(f"  平均置信度: {avg_confidence:.3f}")
        print(f"  高置信度预测 (>0.8): {high_confidence}/{len(metrics['confidences'])} ({high_confidence/len(metrics['confidences'])*100:.1f}%)")
        print()
        
        # 错误分析
        print("🔍 错误分析:")
        false_positives = cm[0, 1]  # 实际负面，预测正面
        false_negatives = cm[1, 0]  # 实际正面，预测负面
        
        print(f"  假正例 (负面→正面): {false_positives} 条")
        print(f"  假负例 (正面→负面): {false_negatives} 条")
        
        if false_positives > false_negatives:
            print("  模型倾向于预测正面情感")
        elif false_negatives > false_positives:
            print("  模型倾向于预测负面情感")
        else:
            print("  模型预测相对平衡")
        
        print()
        print("💡 结论:")
        if metrics['accuracy'] >= 0.9:
            print("  ✅ 模型表现优秀，情感分类能力很强")
        elif metrics['accuracy'] >= 0.8:
            print("  ✅ 模型表现良好，情感分类能力较强")
        elif metrics['accuracy'] >= 0.7:
            print("  ⚠️ 模型表现一般，仍有改进空间")
        else:
            print("  ❌ 模型表现较差，需要进一步训练或调整")
    
    def plot_confusion_matrix(self):
        """绘制混淆矩阵"""
        cm = self.metrics['confusion_matrix']
        
        plt.figure(figsize=(8, 6))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                   xticklabels=['负面', '正面'], 
                   yticklabels=['负面', '正面'])
        plt.title('BERT情感分类混淆矩阵')
        plt.xlabel('预测标签')
        plt.ylabel('真实标签')
        plt.tight_layout()
        plt.savefig('bert_sentiment_confusion_matrix.png', dpi=300, bbox_inches='tight')
        plt.show()
        
        print("✅ 混淆矩阵图已保存到: bert_sentiment_confusion_matrix.png")
    
    def plot_confidence_distribution(self):
        """绘制置信度分布"""
        confidences = self.metrics['confidences']
        predictions = self.metrics['predictions']
        true_labels = self.metrics['true_labels']
        
        # 分离正确和错误的预测
        correct_confidences = [c for c, p, t in zip(confidences, predictions, true_labels) if p == t]
        incorrect_confidences = [c for c, p, t in zip(confidences, predictions, true_labels) if p != t]
        
        plt.figure(figsize=(12, 5))
        
        # 子图1: 整体置信度分布
        plt.subplot(1, 2, 1)
        plt.hist(confidences, bins=20, alpha=0.7, color='blue', edgecolor='black')
        plt.xlabel('置信度')
        plt.ylabel('频次')
        plt.title('整体置信度分布')
        plt.grid(True, alpha=0.3)
        
        # 子图2: 正确vs错误预测的置信度分布
        plt.subplot(1, 2, 2)
        plt.hist(correct_confidences, bins=15, alpha=0.7, label='正确预测', color='green')
        plt.hist(incorrect_confidences, bins=15, alpha=0.7, label='错误预测', color='red')
        plt.xlabel('置信度')
        plt.ylabel('频次')
        plt.title('正确vs错误预测的置信度分布')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig('bert_sentiment_confidence_distribution.png', dpi=300, bbox_inches='tight')
        plt.show()
        
        print("✅ 置信度分布图已保存到: bert_sentiment_confidence_distribution.png")
    
    def save_evaluation_report(self, output_file: str = "bert_sentiment_evaluation_report.json"):
        """保存评估报告"""
        report = {
            'evaluation_time': datetime.now().isoformat(),
            'model_path': self.model_path,
            'tokenizer_path': self.tokenizer_path,
            'device': self.device,
            'metrics': {
                'accuracy': self.metrics['accuracy'],
                'f1_score': self.metrics['f1_score'],
                'precision': self.metrics['precision'],
                'recall': self.metrics['recall'],
                'confusion_matrix': self.metrics['confusion_matrix'].tolist(),
                'avg_inference_time': np.mean(self.metrics['inference_times']),
                'throughput': len(self.metrics['predictions']) / sum(self.metrics['inference_times']),
                'avg_confidence': np.mean(self.metrics['confidences'])
            },
            'sample_predictions': [
                {
                    'text': text[:100] + '...' if len(text) > 100 else text,
                    'true_label': int(label),
                    'predicted_label': int(pred),
                    'confidence': float(conf)
                }
                for text, label, pred, conf in zip(
                    ['Sample text'] * min(10, len(self.metrics['predictions'])),  # 这里需要原始文本
                    self.metrics['true_labels'][:10],
                    self.metrics['predictions'][:10],
                    self.metrics['confidences'][:10]
                )
            ]
        }
        
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(report, f, ensure_ascii=False, indent=2)
        
        print(f"📄 详细评估报告已保存到: {output_file}")

def main():
    parser = argparse.ArgumentParser(description="BERT情感分类模型评估")
    parser.add_argument("--model_path", required=True, help="BERT模型路径")
    parser.add_argument("--tokenizer_path", help="分词器路径（可选，默认使用model_path）")
    parser.add_argument("--data_path", required=True, help="评估数据路径（CSV或JSONL格式）")
    parser.add_argument("--device", default="cuda", help="运行设备")
    parser.add_argument("--output_dir", default="./", help="输出目录")
    parser.add_argument("--plot", action="store_true", help="是否生成图表")
    
    args = parser.parse_args()
    
    # 检查输入文件
    if not os.path.exists(args.model_path):
        print(f"❌ 模型路径不存在: {args.model_path}")
        return
    
    if not os.path.exists(args.data_path):
        print(f"❌ 数据文件不存在: {args.data_path}")
        return
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 创建评估器
    evaluator = BERTSentimentEvaluator(
        model_path=args.model_path,
        tokenizer_path=args.tokenizer_path,
        device=args.device
    )
    
    # 加载数据
    texts, labels = evaluator.load_evaluation_data(args.data_path)
    
    # 运行评估
    metrics = evaluator.evaluate_model(texts, labels)
    
    # 打印结果
    evaluator.print_evaluation_results()
    
    # 生成图表
    if args.plot:
        evaluator.plot_confusion_matrix()
        evaluator.plot_confidence_distribution()
    
    # 保存报告
    output_file = os.path.join(args.output_dir, "bert_sentiment_evaluation_report.json")
    evaluator.save_evaluation_report(output_file)
    
    print("\n🎉 评估完成！")
    print(f"📊 主要指标:")
    print(f"   准确率: {metrics['accuracy']:.4f}")
    print(f"   F1分数: {metrics['f1_score']:.4f}")
    print(f"   平均推理时间: {np.mean(metrics['inference_times']):.4f}s")

if __name__ == "__main__":
    main()
