#!/usr/bin/env python3
"""
Aggregate VLMEvalKit results across checkpoints

把模型目录下所有 checkpoint 的评测结果汇总成 Excel 表格。
- Rows: training steps (checkpoint-500, checkpoint-1000, ...)
- Columns: benchmarks (GQA, MMMU, MathVista, ...)

Usage:
    python aggregate_checkpoints.py --model_dir output_models/j63_mdm_xxx
    python aggregate_checkpoints.py --model_dir output_models/j63_mdm_xxx --output results.xlsx
"""

import argparse
import re
from pathlib import Path
from typing import Dict, Optional

import pandas as pd


# Benchmark 结果文件映射
BENCHMARK_RESULT_PATTERNS = {
    "MathVista": "*MathVista_MINI*score*.csv",
    "MMMU": "*MMMU_DEV_VAL*acc*.csv",
    "MMBench": "*MMBench_DEV_EN_V11*acc*.csv",
    "GQA": "*GQA_TestDev_Balanced*.csv",
    "TextVQA": "*TextVQA_VAL*.csv",
    "POPE": "*POPE*.csv",
}


def find_result_file(eval_dir: Path, pattern: str) -> Optional[Path]:
    """Find result file matching pattern in eval directory."""
    if not eval_dir.exists():
        return None
    
    # Search for result files (also search in subdirectories)
    files = list(eval_dir.glob(pattern))
    files.extend(eval_dir.glob(f"*/{pattern}"))
    files.extend(eval_dir.glob(f"*/*/{pattern}"))
    files.extend(eval_dir.glob(f"*/*/*/{pattern}"))
    
    # Filter out backup directories and symlinks
    real_files = []
    for f in files:
        if "bak_" in str(f):
            continue
        if f.is_symlink():
            continue
        if f.exists():
            real_files.append(f)
    
    if real_files:
        # Use the most recent file
        try:
            return max(real_files, key=lambda f: f.stat().st_mtime)
        except Exception:
            return real_files[0]
    return None


def parse_mathvista_result(file_path: Path) -> Optional[float]:
    """Parse MathVista score CSV."""
    try:
        df = pd.read_csv(file_path)
        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:
        return None


def parse_mmmu_result(file_path: Path) -> Optional[float]:
    """Parse MMMU result CSV."""
    try:
        df = pd.read_csv(file_path)
        if "Overall" in df.columns:
            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
            val = float(df["Overall"].iloc[0])
            return val * 100 if val <= 1 else val
        return None
    except Exception:
        return None


def parse_mmbench_result(file_path: Path) -> Optional[float]:
    """Parse MMBench result CSV."""
    try:
        df = pd.read_csv(file_path)
        if "Overall" in df.columns:
            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
            val = float(df["Overall"].iloc[0])
            return val * 100 if val <= 1 else val
        return None
    except Exception:
        return None


def parse_gqa_result(file_path: Path) -> Optional[float]:
    """Parse GQA result CSV."""
    try:
        df = pd.read_csv(file_path)
        if "Overall" in df.columns:
            return float(df["Overall"].iloc[0])
        return None
    except Exception:
        return None


def parse_textvqa_result(file_path: Path) -> Optional[float]:
    """Parse TextVQA result CSV."""
    try:
        df = pd.read_csv(file_path)
        if "Overall" in df.columns:
            return float(df["Overall"].iloc[0])
        return None
    except Exception:
        return None


def parse_pope_result(file_path: Path) -> Optional[float]:
    """Parse POPE result CSV."""
    try:
        df = pd.read_csv(file_path)
        if "Overall" in df.columns:
            # Find Overall row
            if "split" in df.columns:
                overall_row = df[df["split"] == "Overall"]
                if not overall_row.empty:
                    return float(overall_row["Overall"].iloc[0])
            return float(df["Overall"].iloc[0])
        return None
    except Exception:
        return None


BENCHMARK_PARSERS = {
    "MathVista": parse_mathvista_result,
    "MMMU": parse_mmmu_result,
    "MMBench": parse_mmbench_result,
    "GQA": parse_gqa_result,
    "TextVQA": parse_textvqa_result,
    "POPE": parse_pope_result,
}


def extract_step_number(checkpoint_name: str) -> int:
    """Extract step number from checkpoint name (e.g., checkpoint-1500 -> 1500)."""
    match = re.search(r"checkpoint-(\d+)", checkpoint_name)
    if match:
        return int(match.group(1))
    return 0


def collect_checkpoint_results(model_dir: Path) -> Dict[int, Dict[str, Optional[float]]]:
    """Collect results from all checkpoints in model directory.
    
    Looks for results in:
    - New path: model_dir/eval_vlmeval/checkpoint-XXX/
    - Old path: model_dir/checkpoint-XXX/eval_vlmeval/ (for backwards compatibility)
    """
    results = {}
    
    # Find all checkpoint directories
    checkpoint_dirs = sorted(
        [d for d in model_dir.iterdir() if d.is_dir() and d.name.startswith("checkpoint-")],
        key=lambda d: extract_step_number(d.name)
    )
    
    # Also check model_dir/eval_vlmeval/ for new-style results
    eval_vlmeval_dir = model_dir / "eval_vlmeval"
    if eval_vlmeval_dir.exists():
        for d in eval_vlmeval_dir.iterdir():
            if d.is_dir() and d.name.startswith("checkpoint-"):
                # Add to checkpoint_dirs if not already there
                corresponding_ckpt = model_dir / d.name
                if corresponding_ckpt not in checkpoint_dirs:
                    checkpoint_dirs.append(corresponding_ckpt)
        checkpoint_dirs = sorted(checkpoint_dirs, key=lambda d: extract_step_number(d.name))
    
    print(f"Found {len(checkpoint_dirs)} checkpoints in {model_dir}")
    
    for ckpt_dir in checkpoint_dirs:
        step = extract_step_number(ckpt_dir.name)
        ckpt_name = ckpt_dir.name
        
        # Try new path first: model_dir/eval_vlmeval/checkpoint-XXX/
        eval_dir = model_dir / "eval_vlmeval" / ckpt_name
        
        # Fallback to old path: model_dir/checkpoint-XXX/eval_vlmeval/
        if not eval_dir.exists():
            eval_dir = ckpt_dir / "eval_vlmeval"
        
        if not eval_dir.exists():
            print(f"  {ckpt_name}: no eval_vlmeval directory")
            continue
        
        ckpt_results = {}
        found_any = False
        
        for benchmark, pattern in BENCHMARK_RESULT_PATTERNS.items():
            result_file = find_result_file(eval_dir, pattern)
            if result_file:
                parser = BENCHMARK_PARSERS.get(benchmark)
                if parser:
                    score = parser(result_file)
                    if score is not None:
                        ckpt_results[benchmark] = score
                        found_any = True
        
        if found_any:
            results[step] = ckpt_results
            benchmarks_found = ", ".join(f"{k}: {v:.1f}" for k, v in ckpt_results.items())
            print(f"  {ckpt_name}: {benchmarks_found}")
        else:
            print(f"  {ckpt_name}: no results found")
    
    return results


def create_results_dataframe(results: Dict[int, Dict[str, Optional[float]]]) -> pd.DataFrame:
    """Create DataFrame from results dict."""
    if not results:
        return pd.DataFrame()
    
    # Get all benchmarks
    all_benchmarks = set()
    for step_results in results.values():
        all_benchmarks.update(step_results.keys())
    
    # Sort benchmarks in desired order
    benchmark_order = ["GQA", "MMBench", "MMMU", "MathVista", "POPE", "TextVQA"]
    sorted_benchmarks = [b for b in benchmark_order if b in all_benchmarks]
    sorted_benchmarks.extend([b for b in sorted(all_benchmarks) if b not in benchmark_order])
    
    # Build DataFrame
    rows = []
    for step in sorted(results.keys()):
        row = {"Step": step}
        for benchmark in sorted_benchmarks:
            row[benchmark] = results[step].get(benchmark)
        rows.append(row)
    
    df = pd.DataFrame(rows)
    df = df.set_index("Step")
    
    return df


def main():
    parser = argparse.ArgumentParser(
        description="Aggregate VLMEvalKit results across checkpoints",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--model_dir",
        type=str,
        required=True,
        help="Path to model directory containing checkpoint-* subdirectories",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output Excel file path (default: model_dir/eval_summary.xlsx)",
    )
    
    args = parser.parse_args()
    
    model_dir = Path(args.model_dir)
    if not model_dir.exists():
        print(f"Error: Model directory does not exist: {model_dir}")
        return 1
    
    # Collect results
    results = collect_checkpoint_results(model_dir)
    
    if not results:
        print("\nNo results found.")
        return 1
    
    # Create DataFrame
    df = create_results_dataframe(results)
    
    # Determine output path
    output_path = args.output
    if output_path is None:
        output_path = model_dir / "eval_summary.xlsx"
    else:
        output_path = Path(output_path)
    
    # Save to Excel
    df.to_excel(output_path)
    print(f"\n✓ Results saved to: {output_path}")
    
    # Also print as markdown table
    print("\n" + "=" * 60)
    print(df.round(1).to_markdown())
    print("=" * 60)
    
    return 0


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

