#!/usr/bin/env python3
"""
VLMEvalKit 结果汇总脚本

读取 VLMEvalKit 评测结果并汇总成表格。

Usage:
    python aggregate_vlmeval_results.py --input_dir eval_outputs/vlmeval_all
    python aggregate_vlmeval_results.py --input_dir eval_outputs/vlmeval_all --output summary.md
"""

import argparse
import os
import re
from pathlib import Path
from typing import Dict, List, Optional, Any

import pandas as pd


# Benchmark 结果文件映射
BENCHMARK_RESULT_PATTERNS = {
    "MathVista_MINI": "*MathVista_MINI*score*.csv",
    "MMMU_DEV_VAL": "*MMMU_DEV_VAL*acc*.csv",
    "MMMU_Pro_V_COT": "*MMMU_Pro_V_COT*acc*.csv",
    "MMBench_DEV_EN_V11": "*MMBench_DEV_EN_V11*acc*.csv",
    "GQA_TestDev_Balanced": "*GQA_TestDev_Balanced*.csv",
    "TextVQA_VAL": "*TextVQA_VAL*.csv",
    "POPE": "*POPE*.csv",
    "ChartQA_TEST": "*ChartQA_TEST*.csv",
    "RealWorldQA": "*RealWorldQA*.csv",
    "AI2D_TEST": "*AI2D_TEST*.csv",
    "DocVQA_VAL": "*DocVQA_VAL*.csv",
    "SEEDBench2_Plus": "*SEEDBench2_Plus*.csv",
}

# 各 Benchmark 的主要指标
BENCHMARK_METRICS = {
    "MathVista_MINI": "acc",
    "MMMU_DEV_VAL": "Overall",
    "MMMU_Pro_V_COT": "Overall",
    "MMBench_DEV_EN_V11": "Overall",
    "GQA_TestDev_Balanced": "acc",
    "TextVQA_VAL": "acc",
    "POPE": "Overall",
    "ChartQA_TEST": "Overall",
    "RealWorldQA": "Overall",
    "AI2D_TEST": "Overall",
    "DocVQA_VAL": "Overall",
    "SEEDBench2_Plus": "Overall",
}


def find_result_files(input_dir: Path) -> Dict[str, Path]:
    """Find result files for each benchmark.
    
    VLMEvalKit output structure:
    input_dir/
      MathVista_MINI/
        Qwen2.5-VL-3B-Instruct/
          Qwen2.5-VL-3B-Instruct_MathVista_MINI_gpt-4-0125_score.csv
    """
    results = {}
    
    for benchmark, pattern in BENCHMARK_RESULT_PATTERNS.items():
        benchmark_dir = input_dir / benchmark
        if not benchmark_dir.exists():
            continue
        
        # Search for result files (also search in model subdirectories)
        files = list(benchmark_dir.glob(pattern))
        files.extend(benchmark_dir.glob(f"*/{pattern}"))
        files.extend(benchmark_dir.glob(f"*/*/{pattern}"))
        
        # Filter out backup directories and symlinks (symlinks may be broken)
        real_files = []
        for f in files:
            if "bak_" in str(f):
                continue
            if f.is_symlink():
                continue  # Skip symlinks, use real files only
            if f.exists():
                real_files.append(f)
        
        if real_files:
            # Use the most recent file
            try:
                results[benchmark] = max(real_files, key=lambda f: f.stat().st_mtime)
            except Exception:
                results[benchmark] = real_files[0]
    
    return results


def parse_mathvista_result(file_path: Path) -> Optional[float]:
    """Parse MathVista score CSV."""
    try:
        df = pd.read_csv(file_path)
        # Find Overall accuracy
        if "Task&Skill" in df.columns and "acc" in df.columns:
            overall = df[df["Task&Skill"] == "Overall"]
            if not overall.empty:
                return float(overall["acc"].iloc[0])
        return None
    except Exception as e:
        print(f"Error parsing {file_path}: {e}")
        return None


def parse_mmmu_result(file_path: Path) -> Optional[float]:
    """Parse MMMU result CSV.
    
    Format: "split","Overall",...
            "validation","0.475",...
    """
    try:
        df = pd.read_csv(file_path)
        if "Overall" in df.columns:
            # Find validation split row
            if "split" in df.columns:
                val_row = df[df["split"] == "validation"]
                if not val_row.empty:
                    val = float(val_row["Overall"].iloc[0])
                    return val * 100 if val <= 1 else val
            # Fallback to first row
            val = float(df["Overall"].iloc[0])
            return val * 100 if val <= 1 else val
        return None
    except Exception as e:
        print(f"Error parsing {file_path}: {e}")
        return None


def parse_mmbench_result(file_path: Path) -> Optional[float]:
    """Parse MMBench result CSV.
    
    Format: "split","Overall",...
            "dev","0.768",...
    """
    try:
        df = pd.read_csv(file_path)
        if "Overall" in df.columns:
            # Find dev split row
            if "split" in df.columns:
                dev_row = df[df["split"] == "dev"]
                if not dev_row.empty:
                    val = float(dev_row["Overall"].iloc[0])
                    return val * 100 if val <= 1 else val
            # Fallback to first row
            val = float(df["Overall"].iloc[0])
            return val * 100 if val <= 1 else val
        return None
    except Exception as e:
        print(f"Error parsing {file_path}: {e}")
        return None


def parse_generic_result(file_path: Path, metric: str = "acc") -> Optional[float]:
    """Parse generic VLMEvalKit result CSV."""
    try:
        df = pd.read_csv(file_path)
        
        # Handle split-based CSVs (find "Overall" split row)
        if "split" in df.columns:
            overall_row = df[df["split"] == "Overall"]
            if not overall_row.empty:
                if metric in df.columns:
                    val = float(overall_row[metric].iloc[0])
                    return val * 100 if val <= 1 else val
                elif "Overall" in df.columns:
                    val = float(overall_row["Overall"].iloc[0])
                    return val * 100 if val <= 1 else val
        
        # Direct metric column
        if metric in df.columns:
            val = df[metric].iloc[0]
            return float(val) * 100 if val <= 1 else float(val)
        
        # Overall column
        if "Overall" in df.columns:
            val = df["Overall"].iloc[0]
            return float(val) * 100 if val <= 1 else float(val)
        
        # Try first numeric column
        for col in df.columns:
            try:
                val = float(df[col].iloc[0])
                return val * 100 if val <= 1 else val
            except (ValueError, TypeError):
                continue
        return None
    except Exception as e:
        print(f"Error parsing {file_path}: {e}")
        return None


def parse_result(benchmark: str, file_path: Path) -> Optional[float]:
    """Parse result file based on benchmark type."""
    if benchmark == "MathVista_MINI":
        return parse_mathvista_result(file_path)
    elif benchmark in ["MMMU_DEV_VAL", "MMBench_DEV_EN_V11"]:
        return parse_mmmu_result(file_path)
    else:
        metric = BENCHMARK_METRICS.get(benchmark, "acc")
        return parse_generic_result(file_path, metric)


def collect_results(input_dir: Path) -> Dict[str, float]:
    """Collect all benchmark results from input directory."""
    result_files = find_result_files(input_dir)
    results = {}
    
    for benchmark, file_path in result_files.items():
        score = parse_result(benchmark, file_path)
        if score is not None:
            results[benchmark] = score
            print(f"  {benchmark}: {score:.2f}%")
        else:
            print(f"  {benchmark}: Failed to parse")
    
    return results


def format_markdown_table(
    results: Dict[str, Dict[str, float]],
    benchmarks: Optional[List[str]] = None
) -> str:
    """Format results as markdown table."""
    if not results:
        return "No results found."
    
    # Get all benchmarks
    if benchmarks is None:
        benchmarks = sorted(set(b for r in results.values() for b in r.keys()))
    
    # Build header
    lines = []
    lines.append("# VLMEvalKit 评测结果\n")
    
    header = "| Model |"
    separator = "|-------|"
    for b in benchmarks:
        short_name = b.replace("_MINI", "").replace("_DEV_VAL", "").replace("_DEV_EN_V11", "").replace("_TestDev_Balanced", "").replace("_VAL", "")
        header += f" {short_name} |"
        separator += "--------|"
    
    lines.append(header)
    lines.append(separator)
    
    # Build rows
    for model_name, model_results in results.items():
        row = f"| {model_name} |"
        for b in benchmarks:
            if b in model_results:
                row += f" {model_results[b]:.1f} |"
            else:
                row += " - |"
        lines.append(row)
    
    lines.append("")
    return "\n".join(lines)


def format_plain_table(
    results: Dict[str, Dict[str, float]],
    benchmarks: Optional[List[str]] = None
) -> str:
    """Format results as plain text table."""
    if not results:
        return "No results found."
    
    if benchmarks is None:
        benchmarks = sorted(set(b for r in results.values() for b in r.keys()))
    
    lines = []
    lines.append("=" * 80)
    lines.append("VLMEvalKit 评测结果")
    lines.append("=" * 80)
    
    # Header
    header = f"{'Model':<30}"
    for b in benchmarks:
        short_name = b.replace("_MINI", "").replace("_DEV_VAL", "").replace("_DEV_EN_V11", "").replace("_TestDev_Balanced", "").replace("_VAL", "")
        header += f" {short_name:>12}"
    lines.append(header)
    lines.append("-" * 80)
    
    # Rows
    for model_name, model_results in results.items():
        row = f"{model_name:<30}"
        for b in benchmarks:
            if b in model_results:
                row += f" {model_results[b]:>11.1f}%"
            else:
                row += f" {'-':>12}"
        lines.append(row)
    
    lines.append("=" * 80)
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="Aggregate VLMEvalKit results")
    parser.add_argument(
        "--input_dir",
        type=str,
        required=True,
        help="Input directory containing benchmark results"
    )
    parser.add_argument(
        "--model_name",
        type=str,
        default=None,
        help="Model name for display (default: inferred from directory)"
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output file path (default: print to stdout)"
    )
    parser.add_argument(
        "--format",
        type=str,
        choices=["markdown", "plain"],
        default="markdown",
        help="Output format"
    )
    parser.add_argument(
        "--compare_dirs",
        type=str,
        nargs="+",
        default=None,
        help="Additional directories to compare (for multi-model comparison)"
    )
    
    args = parser.parse_args()
    
    input_dir = Path(args.input_dir)
    if not input_dir.exists():
        print(f"Error: Input directory does not exist: {input_dir}")
        return 1
    
    # Collect results
    all_results = {}
    
    # Main input directory
    model_name = args.model_name or input_dir.name
    print(f"Collecting results for: {model_name}")
    results = collect_results(input_dir)
    if results:
        all_results[model_name] = results
    
    # Compare directories
    if args.compare_dirs:
        for compare_dir in args.compare_dirs:
            compare_path = Path(compare_dir)
            if compare_path.exists():
                compare_name = compare_path.name
                print(f"\nCollecting results for: {compare_name}")
                compare_results = collect_results(compare_path)
                if compare_results:
                    all_results[compare_name] = compare_results
    
    if not all_results:
        print("No results found.")
        return 1
    
    # Format output
    if args.format == "markdown":
        output = format_markdown_table(all_results)
    else:
        output = format_plain_table(all_results)
    
    # Output
    if args.output:
        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "w") as f:
            f.write(output)
        print(f"\nResults saved to: {output_path}")
    else:
        print("\n" + output)
    
    return 0


if __name__ == "__main__":
    exit(main())


