#!/usr/bin/env python3
"""
多配置评测结果汇总脚本

读取 VLMEvalKit 多配置评测结果并生成对比表格。

支持的目录结构:
    output_models/exp_name/eval_vlmeval/
    ├── MDM=0_SPEC=0/
    │   ├── checkpoint-1000/
    │   │   └── summary.md
    │   └── checkpoint-2000/
    ├── MDM=1_SPEC=0/
    │   └── checkpoint-*/
    └── MDM=1_SPEC=1/
        └── checkpoint-*/

Usage:
    python aggregate_multi_config_results.py output_models/xxx/eval_vlmeval
    python aggregate_multi_config_results.py output_models/xxx/eval_vlmeval --output comparison.md
    python aggregate_multi_config_results.py output_models/xxx/eval_vlmeval --best-checkpoint
"""

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

# Add parent directory to path for imports
_script_dir = os.path.dirname(os.path.abspath(__file__))
_vlmeval_dir = os.path.join(_script_dir, "vlmeval")
if _vlmeval_dir not in sys.path:
    sys.path.insert(0, _vlmeval_dir)

try:
    import pandas as pd
except ImportError:
    pd = None

# Import from aggregate_vlmeval_results
try:
    from aggregate_vlmeval_results import (
        find_result_files,
        parse_result,
        BENCHMARK_RESULT_PATTERNS,
        BENCHMARK_METRICS,
    )
except ImportError:
    # Fallback if import fails
    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_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",
    }
    find_result_files = None
    parse_result = None


# Config name descriptions
CONFIG_DESCRIPTIONS = {
    "MDM=0_SPEC=0": "AR (Autoregressive)",
    "MDM=1_SPEC=0": "MDM Standard",
    "MDM=1_SPEC=1": "MDM Speculative",
}


def parse_config_name(config_dir: str) -> Tuple[int, int]:
    """Parse config directory name to (USE_MDM, USE_SPECULATIVE) tuple."""
    match = re.match(r"MDM=(\d)_SPEC=(\d)", config_dir)
    if match:
        return int(match.group(1)), int(match.group(2))
    return None, None


def extract_checkpoint_step(ckpt_name: str) -> Optional[int]:
    """Extract step number from checkpoint name."""
    match = re.match(r"checkpoint-(\d+)", ckpt_name)
    if match:
        return int(match.group(1))
    return None


def find_config_dirs(eval_dir: Path) -> List[str]:
    """Find all configuration directories (MDM=X_SPEC=Y format)."""
    configs = []
    for d in eval_dir.iterdir():
        if d.is_dir() and re.match(r"MDM=\d_SPEC=\d", d.name):
            configs.append(d.name)
    return sorted(configs)


def find_checkpoint_dirs(config_dir: Path) -> List[str]:
    """Find all checkpoint directories in a config directory."""
    checkpoints = []
    for d in config_dir.iterdir():
        if d.is_dir() and d.name.startswith("checkpoint-"):
            checkpoints.append(d.name)
    # Sort by step number
    return sorted(checkpoints, key=lambda x: extract_checkpoint_step(x) or 0)


def collect_checkpoint_results(checkpoint_dir: Path) -> Dict[str, float]:
    """Collect benchmark results from a checkpoint directory."""
    results = {}
    
    for benchmark in BENCHMARK_RESULT_PATTERNS.keys():
        benchmark_dir = checkpoint_dir / benchmark
        if not benchmark_dir.exists():
            continue
        
        pattern = BENCHMARK_RESULT_PATTERNS[benchmark]
        files = list(benchmark_dir.glob(pattern))
        files.extend(benchmark_dir.glob(f"*/{pattern}"))
        files.extend(benchmark_dir.glob(f"*/*/{pattern}"))
        
        # Filter valid files
        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:
            try:
                file_path = max(real_files, key=lambda f: f.stat().st_mtime)
            except Exception:
                file_path = real_files[0]
            
            score = parse_result_file(benchmark, file_path)
            if score is not None:
                results[benchmark] = score
    
    return results


def parse_result_file(benchmark: str, file_path: Path) -> Optional[float]:
    """Parse a benchmark result file."""
    if pd is None:
        return None
    
    try:
        df = pd.read_csv(file_path)
        
        if benchmark == "MathVista_MINI":
            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])
        
        elif benchmark in ["MMMU_DEV_VAL", "MMBench_DEV_EN_V11"]:
            if "Overall" in df.columns:
                if "split" in df.columns:
                    for split in ["validation", "dev"]:
                        split_row = df[df["split"] == split]
                        if not split_row.empty:
                            val = float(split_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
        
        else:
            metric = BENCHMARK_METRICS.get(benchmark, "acc")
            
            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
            
            if metric in df.columns:
                val = df[metric].iloc[0]
                return float(val) * 100 if val <= 1 else float(val)
            
            if "Overall" in df.columns:
                val = df["Overall"].iloc[0]
                return float(val) * 100 if val <= 1 else float(val)
            
            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 collect_all_results(eval_dir: Path) -> Dict[str, Dict[str, Dict[str, float]]]:
    """Collect all results organized by config -> checkpoint -> benchmark.
    
    Returns:
        {config: {checkpoint: {benchmark: score}}}
    """
    all_results = {}
    
    config_dirs = find_config_dirs(eval_dir)
    
    for config in config_dirs:
        config_path = eval_dir / config
        all_results[config] = {}
        
        checkpoint_dirs = find_checkpoint_dirs(config_path)
        
        for checkpoint in checkpoint_dirs:
            checkpoint_path = config_path / checkpoint
            results = collect_checkpoint_results(checkpoint_path)
            if results:
                all_results[config][checkpoint] = results
    
    return all_results


def compute_average(results: Dict[str, float], benchmarks: Optional[List[str]] = None) -> float:
    """Compute average score across benchmarks."""
    if not results:
        return 0.0
    
    if benchmarks:
        scores = [results.get(b) for b in benchmarks if results.get(b) is not None]
    else:
        scores = list(results.values())
    
    return sum(scores) / len(scores) if scores else 0.0


def format_comparison_table(
    all_results: Dict[str, Dict[str, Dict[str, float]]],
    benchmarks: Optional[List[str]] = None,
    best_checkpoint_only: bool = False,
) -> str:
    """Format results as a markdown comparison table.
    
    Args:
        all_results: {config: {checkpoint: {benchmark: score}}}
        benchmarks: List of benchmarks to include (None = all)
        best_checkpoint_only: If True, only show best checkpoint per config
    """
    if not all_results:
        return "No results found.\n"
    
    # Get all benchmarks
    all_benchmarks = set()
    for config_results in all_results.values():
        for ckpt_results in config_results.values():
            all_benchmarks.update(ckpt_results.keys())
    
    if benchmarks is None:
        benchmarks = sorted(all_benchmarks)
    
    lines = []
    lines.append("# Multi-Configuration Evaluation Results\n")
    
    # Summary section
    lines.append("## Configuration Summary\n")
    lines.append("| Config | Description | Checkpoints |")
    lines.append("|--------|-------------|-------------|")
    for config in sorted(all_results.keys()):
        desc = CONFIG_DESCRIPTIONS.get(config, config)
        num_ckpts = len(all_results[config])
        lines.append(f"| {config} | {desc} | {num_ckpts} |")
    lines.append("")
    
    # Detailed comparison per checkpoint
    if best_checkpoint_only:
        lines.append("## Best Checkpoint per Configuration\n")
    else:
        lines.append("## Detailed Results\n")
    
    # Short benchmark names
    short_names = {}
    for b in benchmarks:
        short = b.replace("_MINI", "").replace("_DEV_VAL", "").replace("_DEV_EN_V11", "")
        short = short.replace("_TestDev_Balanced", "").replace("_VAL", "").replace("_TEST", "")
        short_names[b] = short
    
    # Header
    header = "| Config | Checkpoint |"
    separator = "|--------|------------|"
    for b in benchmarks:
        header += f" {short_names[b]} |"
        separator += "--------|"
    header += " Avg |"
    separator += "------|"
    
    lines.append(header)
    lines.append(separator)
    
    # Rows
    for config in sorted(all_results.keys()):
        config_results = all_results[config]
        
        if best_checkpoint_only:
            # Find checkpoint with best average
            best_ckpt = None
            best_avg = -1
            for ckpt, results in config_results.items():
                avg = compute_average(results, benchmarks)
                if avg > best_avg:
                    best_avg = avg
                    best_ckpt = ckpt
            
            if best_ckpt:
                results = config_results[best_ckpt]
                row = f"| {config} | {best_ckpt} |"
                for b in benchmarks:
                    if b in results:
                        row += f" {results[b]:.1f} |"
                    else:
                        row += " - |"
                avg = compute_average(results, benchmarks)
                row += f" {avg:.1f} |"
                lines.append(row)
        else:
            # Show all checkpoints
            for ckpt in sorted(config_results.keys(), key=lambda x: extract_checkpoint_step(x) or 0):
                results = config_results[ckpt]
                row = f"| {config} | {ckpt} |"
                for b in benchmarks:
                    if b in results:
                        row += f" {results[b]:.1f} |"
                    else:
                        row += " - |"
                avg = compute_average(results, benchmarks)
                row += f" {avg:.1f} |"
                lines.append(row)
    
    lines.append("")
    
    # Per-config best comparison
    if not best_checkpoint_only and len(all_results) > 1:
        lines.append("## Best per Config Comparison\n")
        
        header = "| Config |"
        separator = "|--------|"
        for b in benchmarks:
            header += f" {short_names[b]} |"
            separator += "--------|"
        header += " Avg |"
        separator += "------|"
        
        lines.append(header)
        lines.append(separator)
        
        for config in sorted(all_results.keys()):
            config_results = all_results[config]
            
            # Find best for each benchmark
            best_scores = {}
            for results in config_results.values():
                for b, score in results.items():
                    if b not in best_scores or score > best_scores[b]:
                        best_scores[b] = score
            
            row = f"| {config} |"
            for b in benchmarks:
                if b in best_scores:
                    row += f" {best_scores[b]:.1f} |"
                else:
                    row += " - |"
            avg = compute_average(best_scores, benchmarks)
            row += f" {avg:.1f} |"
            lines.append(row)
        
        lines.append("")
    
    return "\n".join(lines)


def format_csv_output(
    all_results: Dict[str, Dict[str, Dict[str, float]]],
    benchmarks: Optional[List[str]] = None,
) -> str:
    """Format results as CSV for easy import."""
    if not all_results:
        return ""
    
    all_benchmarks = set()
    for config_results in all_results.values():
        for ckpt_results in config_results.values():
            all_benchmarks.update(ckpt_results.keys())
    
    if benchmarks is None:
        benchmarks = sorted(all_benchmarks)
    
    lines = []
    
    # Header
    header = "Config,Checkpoint," + ",".join(benchmarks) + ",Avg"
    lines.append(header)
    
    # Rows
    for config in sorted(all_results.keys()):
        config_results = all_results[config]
        for ckpt in sorted(config_results.keys(), key=lambda x: extract_checkpoint_step(x) or 0):
            results = config_results[ckpt]
            row_values = [config, ckpt]
            for b in benchmarks:
                if b in results:
                    row_values.append(f"{results[b]:.2f}")
                else:
                    row_values.append("")
            avg = compute_average(results, benchmarks)
            row_values.append(f"{avg:.2f}")
            lines.append(",".join(row_values))
    
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(
        description="Aggregate multi-configuration VLMEvalKit results",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Show all results
  python aggregate_multi_config_results.py output_models/xxx/eval_vlmeval
  
  # Only show best checkpoint per config
  python aggregate_multi_config_results.py output_models/xxx/eval_vlmeval --best-checkpoint
  
  # Save to file
  python aggregate_multi_config_results.py output_models/xxx/eval_vlmeval --output comparison.md
  
  # Export as CSV
  python aggregate_multi_config_results.py output_models/xxx/eval_vlmeval --format csv
"""
    )
    parser.add_argument(
        "eval_dir",
        type=str,
        help="Eval directory containing config subdirectories (MDM=X_SPEC=Y/checkpoint-*)"
    )
    parser.add_argument(
        "--output", "-o",
        type=str,
        default=None,
        help="Output file path (default: print to stdout)"
    )
    parser.add_argument(
        "--format", "-f",
        type=str,
        choices=["markdown", "csv"],
        default="markdown",
        help="Output format (default: markdown)"
    )
    parser.add_argument(
        "--best-checkpoint",
        action="store_true",
        help="Only show best checkpoint per configuration"
    )
    parser.add_argument(
        "--benchmarks",
        type=str,
        nargs="+",
        default=None,
        help="Specific benchmarks to include (default: all)"
    )
    
    args = parser.parse_args()
    
    eval_dir = Path(args.eval_dir)
    if not eval_dir.exists():
        print(f"Error: Directory does not exist: {eval_dir}")
        return 1
    
    print(f"Collecting results from: {eval_dir}")
    
    # Check if this is a multi-config directory
    config_dirs = find_config_dirs(eval_dir)
    
    if not config_dirs:
        # Might be a single-config legacy directory
        print("No multi-config directories found (MDM=X_SPEC=Y).")
        print("This script is for multi-config evaluation results.")
        print("For single-config results, use: aggregate_vlmeval_results.py")
        return 1
    
    print(f"Found {len(config_dirs)} configurations: {', '.join(config_dirs)}")
    
    # Collect all results
    all_results = collect_all_results(eval_dir)
    
    total_checkpoints = sum(len(ckpts) for ckpts in all_results.values())
    print(f"Collected results from {total_checkpoints} checkpoints")
    
    # Format output
    if args.format == "csv":
        output = format_csv_output(all_results, args.benchmarks)
    else:
        output = format_comparison_table(
            all_results,
            args.benchmarks,
            args.best_checkpoint
        )
    
    # 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())

