import os
from multimetric_eval import ModelEvaluator, load_dataset

# ================= 配置区域 =================

# 1. 模型路径配置
MODEL_PATHS = {
    # [修改点] 这里改为 PyTorch 版本的本地模型路径
    "bleurt": "./model/lucadiliello/BLEURT-20", 
    "comet": "Unbabel/wmt22-comet-da"
}

# 2. 手动输入的翻译结果列表 (Hypothesis)
manual_hypotheses = [
    "To Leon Werth",  # 第1句
    "I apologize to the children who read this book for dedicating it to an adult.", # 第2句
    "I have a serious excuse: this grown-up is the best friend I have in the world.", # 第3句
    "I have another excuse: this grown-up understands everything, even books for children.", # 第4句
    "I have a third excuse: this grown-up lives in France where he is hungry and cold." # 第5句
]

# ===========================================

def main():
    # 1. 加载内置数据集 (获取 Reference 和 Source)
    print("📦 加载数据集 (zh-en-littleprince)...")
    dataset = load_dataset("zh-en-littleprince")
    
    # 2. 数据对齐
    # 因为我们手动输入的列表可能只有几句，所以需要把数据集切片，使其长度一致
    num_samples = len(manual_hypotheses)
    
    if num_samples > len(dataset):
        print("⚠️ 警告: 手动输入的句子数量超过了数据集总数！")
        return

    print(f"📄 待评测样本数: {num_samples}")
    
    # 获取对应的参考译文和源文本
    reference_texts = dataset.reference_texts[:num_samples]
    source_texts = dataset.source_texts[:num_samples]

    # 3. 初始化评测器 (PyTorch Version)
    print("\n🔧 初始化评测器 (PyTorch Version / Text Only)...")
    
    # 检查本地 PyTorch BLEURT 路径是否存在
    use_bleurt = os.path.exists(MODEL_PATHS['bleurt'])
    if not use_bleurt:
        print(f"⚠️ 找不到 BLEURT 路径: {MODEL_PATHS['bleurt']}")
        print(f"   将在评测中跳过 BLEURT 指标")
    else:
        print(f"✅ 发现本地 BLEURT 模型: {MODEL_PATHS['bleurt']}")

    evaluator = ModelEvaluator(
        use_comet=True,
        use_bleurt=use_bleurt,
        use_whisper=False,      # ❌ 关闭 ASR
        comet_model=MODEL_PATHS['comet'],
        bleurt_path=MODEL_PATHS['bleurt'] # 传入本地路径，ModelEvaluator会自动用 PyTorch 方式加载
    )

    # 4. 运行评测
    print("\n📊 开始计算文本指标 (BLEU, chrF++, COMET, BLEURT)...")
    
    # 这里我们直接调用 evaluate_all，只传入 target_text
    results = evaluator.evaluate_all(
        reference=reference_texts,
        source=source_texts,
        target_text=manual_hypotheses,  # 传入你的列表
        target_speech=None              # ❌ 不传音频
    )

    # 5. 打印结果
    print("\n" + "="*40)
    print("✅ 评测结果报告")
    print("="*40)
    
    # 打印具体的指标得分
    for k, v in results.items():
        if not isinstance(v, list): # 跳过列表类型的详细数据，只打印分数
            print(f"{k:<15}: {v}")
            
    print("-" * 40)
    
    # (可选) 打印对照查看
    print("🔍 抽样对照 (Hypothesis vs Reference):")
    for i in range(min(3, num_samples)):
        print(f"\n[Sample {i}]")
        print(f"  Hyp: {manual_hypotheses[i]}")
        print(f"  Ref: {reference_texts[i]}")
    
    print("="*40)

    # 清理显存
    evaluator.cleanup()

if __name__ == "__main__":
    main()