import argparse
import csv
import json
import os
import re
from pathlib import Path


DEFAULT_FILE = (
    "/225040511/project/Biomni-ReAct/LAB-bench/"
    "cloningscenarios_cloning_20260524_163903/cloningscenarios_results.jsonl"
)


def normalize_choice(value):
    """Return the first A-E option from a model/gold answer string."""
    if value is None:
        return ""
    text = str(value).strip().upper()
    if text in {"A", "B", "C", "D", "E"}:
        return text
    match = re.search(r"\b([A-E])\b", text)
    return match.group(1) if match else text


def compute_accuracy_from_jsonl(file_path, print_errors=False, dedup=True):
    """
    从 JSONL 文件中读取数据，计算 agent_answer 与 answer 的准确率。
    
    Args:
        file_path (str): JSONL 文件路径。
        print_errors (bool): 是否打印错误条目。
        dedup (bool): 是否按 task_id/id 去重，保留最后一次结果。
    
    Returns:
        dict: accuracy/count/error metrics.
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件不存在: {file_path}")

    records = []
    seen = {}
    duplicate_count = 0

    with open(file_path, "r", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                data = json.loads(line)
            except json.JSONDecodeError as e:
                print(f"警告: 第 {line_num} 行 JSON 解析错误: {e}，跳过")
                continue

            data["_line_num"] = line_num
            key = data.get("task_id") or data.get("id")
            if dedup and key:
                if key in seen:
                    duplicate_count += 1
                seen[key] = data
            else:
                records.append(data)

    if dedup:
        records = list(seen.values()) + records

    correct = 0
    total = 0
    missing_answer = 0
    error_count = 0
    error_records = []

    for data in records:
        line_num = data.get("_line_num", "N/A")
        answer = normalize_choice(data.get("answer"))
        agent_answer = normalize_choice(data.get("agent_answer"))
        if not answer:
            print(f"警告: 第 {line_num} 行缺少 'answer' 字段，跳过")
            continue

        total += 1
        if not agent_answer:
            missing_answer += 1
        if data.get("error"):
            error_count += 1

        if answer == agent_answer:
            correct += 1
        else:
            error_records.append(
                {
                    "line_num": line_num,
                    "task_id": data.get("task_id", data.get("id", "N/A")),
                    "answer": answer,
                    "agent_answer": agent_answer,
                    "error": data.get("error", ""),
                    "question": data.get("question", "N/A"),
                }
            )

    accuracy = correct / total if total > 0 else 0.0
    metrics = {
        "file": str(file_path),
        "total": total,
        "correct": correct,
        "wrong": total - correct,
        "accuracy": accuracy,
        "accuracy_percent": accuracy * 100,
        "missing_agent_answer": missing_answer,
        "error_count": error_count,
        "duplicate_count": duplicate_count,
        "error_records": error_records,
    }
    
    if print_errors and error_records:
        print("\n" + "="*80)
        print("错误条目列表 (answer != agent_answer):")
        print("="*80)
        for idx, err in enumerate(error_records, 1):
            print(f"[{idx:03d}] 行号: {err['line_num']}")
            print(f"  Task ID: {err['task_id']}")
            print(f"  问题: {err['question'][:100]}..." if len(err['question']) > 100 else f"  问题: {err['question']}")
            print(f"  正确答案: {err['answer']}")
            print(f"  Agent答案: {err['agent_answer']}")
            if err["error"]:
                print(f"  Error: {err['error']}")
            print("-"*80)
    
    return metrics


def print_metrics(metrics):
    print(f"\n文件: {metrics['file']}")
    print(f"总样本数: {metrics['total']}")
    print(f"正确数: {metrics['correct']}")
    print(f"错误数: {metrics['wrong']}")
    print(f"准确率: {metrics['accuracy']:.4f} ({metrics['accuracy_percent']:.2f}%)")
    print(f"空答案数: {metrics['missing_agent_answer']}")
    print(f"带 error 字段数: {metrics['error_count']}")
    print(f"重复 task_id 行数: {metrics['duplicate_count']}")


def write_error_csv(metrics, out_path):
    fields = ["line_num", "task_id", "answer", "agent_answer", "error", "question"]
    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with out_path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(metrics["error_records"])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Calculate LAB-Bench MCQ accuracy from JSONL results.")
    parser.add_argument("file", nargs="?", default=DEFAULT_FILE, help="Path to *_results.jsonl.")
    parser.add_argument("--print-errors", action="store_true", help="Print wrong records.")
    parser.add_argument("--no-dedup", action="store_true", help="Do not deduplicate by task_id/id.")
    parser.add_argument("--error-csv", default="", help="Optional CSV path for wrong records.")
    args = parser.parse_args()

    try:
        metrics = compute_accuracy_from_jsonl(
            args.file,
            print_errors=args.print_errors,
            dedup=not args.no_dedup,
        )
        print_metrics(metrics)
        if args.error_csv:
            write_error_csv(metrics, args.error_csv)
            print(f"\n错误样本 CSV 已保存: {args.error_csv}")
    except FileNotFoundError as e:
        print(f"错误: {e}")
    except Exception as e:
        print(f"发生未预期的错误: {e}")