#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
汇总多个 checkpoint 的评测结果并画图。

从 summary.md 文件中提取各 benchmark 的分数，汇总成表格并绘制折线图。

Usage:
    python plot_checkpoint_results.py --input_dir VLMEvalKit_local/output_models/xxx
    python plot_checkpoint_results.py --input_dir VLMEvalKit_local/output_models/xxx --benchmarks "GQA,POPE,TextVQA"
"""

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

import matplotlib.pyplot as plt
import pandas as pd


# Known benchmark directories
KNOWN_BENCHMARKS = [
    "GQA_TestDev_Balanced",
    "MMMU_DEV_VAL",
    "MMMU_Pro_V_COT",
    "MMBench_DEV_EN_V11",
    "MathVista_MINI",
    "POPE",
    "TextVQA_VAL",
    "ChartQA_TEST",
    "RealWorldQA",
    "AI2D_TEST",
    "DocVQA_VAL",
    "SEEDBench2_Plus",
]


# 尝试使用 seaborn 风格
try:
    plt.style.use('seaborn-v0_8-whitegrid')
except:
    try:
        plt.style.use('seaborn-whitegrid')
    except:
        pass


def generate_missing_summary(ckpt_dir: Path) -> bool:
    """
    为缺失 summary.md 的 checkpoint 生成汇总。
    
    Args:
        ckpt_dir: checkpoint 目录 (如 eval_vlmeval/checkpoint-1000)
    
    Returns:
        True if summary.md exists or was generated successfully
    """
    summary_path = ckpt_dir / "summary.md"
    
    # Already exists
    if summary_path.exists():
        return True
    
    # Check if has any benchmark results
    has_results = any(
        (ckpt_dir / benchmark).exists() and (ckpt_dir / benchmark).is_dir()
        for benchmark in KNOWN_BENCHMARKS
    )
    
    if not has_results:
        return False
    
    # Find aggregate_vlmeval_results.py script
    script_dir = Path(__file__).parent
    aggregate_script = script_dir / "vlmeval" / "aggregate_vlmeval_results.py"
    
    if not aggregate_script.exists():
        print(f"  Warning: aggregate_vlmeval_results.py not found at {aggregate_script}")
        return False
    
    # Call aggregate_vlmeval_results.py to generate summary.md
    print(f"  Generating missing summary.md for {ckpt_dir.name}...")
    try:
        result = subprocess.run(
            [
                sys.executable,
                str(aggregate_script),
                "--input_dir", str(ckpt_dir),
                "--output", str(summary_path),
            ],
            capture_output=True,
            text=True,
            timeout=60,
        )
        if result.returncode != 0:
            print(f"    Warning: Failed to generate summary: {result.stderr[:200]}")
            return False
        return summary_path.exists()
    except subprocess.TimeoutExpired:
        print(f"    Warning: Timeout generating summary for {ckpt_dir.name}")
        return False
    except Exception as e:
        print(f"    Warning: Error generating summary: {e}")
        return False


def parse_summary_md(summary_path: Path) -> Dict[str, float]:
    """
    解析 summary.md 文件，提取各 benchmark 的分数。
    
    Returns:
        Dict[benchmark_name, score]
    """
    if not summary_path.exists():
        return {}
    
    content = summary_path.read_text()
    lines = content.strip().split('\n')
    
    # 找到表头和数据行
    header_line = None
    data_line = None
    
    for line in lines:
        if line.startswith('| Model'):
            header_line = line
        elif line.startswith('| ') and 'Model' not in line and '---' not in line:
            data_line = line
    
    if not header_line or not data_line:
        return {}
    
    # 解析表头 (去掉 Model 列)
    headers = [h.strip() for h in header_line.split('|')[2:-1]]  # Skip first empty, Model column, and last empty
    
    # 解析数据 (去掉 Model 列)
    values = [v.strip() for v in data_line.split('|')[2:-1]]
    
    results = {}
    for header, value in zip(headers, values):
        try:
            results[header] = float(value)
        except ValueError:
            # 跳过无法解析的值 (如 '-')
            pass
    
    return results


def extract_checkpoint_step(checkpoint_name: str) -> int:
    """从 checkpoint 名称中提取 step 数。"""
    match = re.search(r'checkpoint-(\d+)', checkpoint_name)
    if match:
        return int(match.group(1))
    return 0


def collect_tokens_per_step_for_checkpoint(ckpt_dir: Path) -> Optional[Dict]:
    """
    收集单个 checkpoint 下所有任务（benchmark 子目录）的 tokens_per_step_stats.json，
    并计算所有任务的平均信息。

    Returns:
        包含 avg_mean, avg_std, num_tasks, per_task 等键的字典；若无任何 stats 则返回 None。
    """
    aggregated = []
    per_task = {}

    for subdir in ckpt_dir.iterdir():
        if not subdir.is_dir():
            continue
        stats_file = subdir / "tokens_per_step_stats.json"
        if not stats_file.exists():
            continue
        try:
            data = json.loads(stats_file.read_text())
        except (json.JSONDecodeError, OSError):
            continue
        if not isinstance(data, dict) or "mean" not in data:
            continue
        aggregated.append(data)
        per_task[subdir.name] = {
            "mean": data["mean"],
            "std": data.get("std"),
            "tokens_per_step_ratio": data.get("tokens_per_step_ratio"),
            "total_tokens": data.get("total_tokens"),
            "total_steps": data.get("total_steps"),
            "total_samples": data.get("total_samples"),
        }

    if not aggregated:
        return None

    # 所有任务的平均
    n = len(aggregated)
    avg_mean = sum(d["mean"] for d in aggregated) / n
    avg_std = sum(d.get("std") or 0 for d in aggregated) / n
    avg_ratio = sum(d.get("tokens_per_step_ratio") or d["mean"] for d in aggregated) / n
    sum_total_tokens = sum(d.get("total_tokens") or 0 for d in aggregated)
    sum_total_steps = sum(d.get("total_steps") or 0 for d in aggregated)
    sum_total_samples = sum(d.get("total_samples") or 0 for d in aggregated)

    return {
        "num_tasks": n,
        "avg_mean": avg_mean,
        "avg_std": avg_std,
        "avg_tokens_per_step_ratio": avg_ratio,
        "sum_total_tokens": sum_total_tokens,
        "sum_total_steps": sum_total_steps,
        "sum_total_samples": sum_total_samples,
        "per_task": per_task,
    }


def collect_all_results(input_dir: Path) -> Tuple[pd.DataFrame, List[str], List[dict]]:
    """
    收集所有 checkpoint 的评测结果，以及每个 checkpoint 所有任务的 tokens_per_step 平均信息。

    Args:
        input_dir: 包含 eval_vlmeval/checkpoint-*/summary.md 的目录

    Returns:
        (DataFrame with columns [step, benchmark1, ..., tokens_per_step_avg_mean, ...],
         list of benchmark names,
         list of per-checkpoint tokens_per_step 汇总 dict，用于保存 JSON)
    """
    eval_dir = input_dir
    if not eval_dir.exists():
        print(f"Error: eval_vlmeval directory not found at {eval_dir}")
        return pd.DataFrame(), [], []

    # 收集所有 checkpoint 结果
    all_results = []
    tps_per_checkpoint: List[dict] = []
    missing_count = 0
    generated_count = 0

    for ckpt_dir in sorted(eval_dir.iterdir()):
        if not ckpt_dir.is_dir() or not ckpt_dir.name.startswith('checkpoint-'):
            continue

        summary_path = ckpt_dir / "summary.md"

        # Try to generate missing summary.md
        if not summary_path.exists():
            missing_count += 1
            if generate_missing_summary(ckpt_dir):
                generated_count += 1

        results = parse_summary_md(summary_path)

        # 收集该 checkpoint 下所有任务的 tokens_per_step 平均信息
        tps = collect_tokens_per_step_for_checkpoint(ckpt_dir)
        if tps:
            results["tokens_per_step_avg_mean"] = tps["avg_mean"]
            results["tokens_per_step_avg_std"] = tps["avg_std"]
            results["tokens_per_step_num_tasks"] = tps["num_tasks"]
            step = extract_checkpoint_step(ckpt_dir.name)
            tps_per_checkpoint.append({
                "checkpoint": ckpt_dir.name,
                "step": step,
                "num_tasks": tps["num_tasks"],
                "avg_mean": tps["avg_mean"],
                "avg_std": tps["avg_std"],
                "avg_tokens_per_step_ratio": tps["avg_tokens_per_step_ratio"],
                "sum_total_tokens": tps["sum_total_tokens"],
                "sum_total_steps": tps["sum_total_steps"],
                "sum_total_samples": tps["sum_total_samples"],
                "per_task": tps["per_task"],
            })

        if results:
            step = extract_checkpoint_step(ckpt_dir.name)
            results["step"] = step
            results["checkpoint"] = ckpt_dir.name
            all_results.append(results)
            print(f"  Loaded {ckpt_dir.name}: {len(results)-2} benchmarks" + (
                f", tokens_per_step_avg_mean={results.get('tokens_per_step_avg_mean', 0):.3f}" if tps else ""
            ))
    
    if missing_count > 0:
        print(f"\n  Summary: {missing_count} checkpoints were missing summary.md, {generated_count} generated")
    
    if not all_results:
        print("No results found.")
        return pd.DataFrame(), [], []

    # 转换为 DataFrame
    df = pd.DataFrame(all_results)
    df = df.sort_values("step").reset_index(drop=True)

    # 获取所有 benchmark 名称 (排除 step、checkpoint 和 tokens_per_step 相关列)
    tps_cols = {"tokens_per_step_avg_mean", "tokens_per_step_avg_std", "tokens_per_step_num_tasks"}
    benchmarks = [col for col in df.columns if col not in ["step", "checkpoint"] and col not in tps_cols]

    return df, benchmarks, tps_per_checkpoint


def plot_benchmark_curves(
    df: pd.DataFrame,
    benchmarks: List[str],
    output_path: Path,
    title: str = "Benchmark Scores vs Training Steps",
    figsize: Tuple[int, int] = (14, 10),
):
    """
    绘制各 benchmark 随 training steps 变化的折线图。
    """
    if df.empty:
        print("No data to plot.")
        return
    
    # 过滤存在的 benchmarks
    available_benchmarks = [b for b in benchmarks if b in df.columns]
    if not available_benchmarks:
        print("No matching benchmarks found.")
        return
    
    # 创建图形
    fig, ax = plt.subplots(figsize=figsize)
    
    # 配色方案 - 使用区分度高的颜色
    colors = [
        '#e6194B',  # Red
        '#3cb44b',  # Green
        '#4363d8',  # Blue
        '#f58231',  # Orange
        '#911eb4',  # Purple
        '#42d4f4',  # Cyan
        '#f032e6',  # Magenta
        '#bfef45',  # Lime
        '#fabed4',  # Pink
        '#469990',  # Teal
        '#dcbeff',  # Lavender
        '#9A6324',  # Brown
    ]
    
    # 绘制每个 benchmark
    for i, benchmark in enumerate(available_benchmarks):
        color = colors[i % len(colors)]
        ax.plot(
            df['step'],
            df[benchmark],
            marker='o',
            markersize=4,
            linewidth=2,
            label=benchmark,
            color=color,
        )
    
    # 设置图形属性
    ax.set_xlabel('Training Steps', fontsize=12)
    ax.set_ylabel('Score (%)', fontsize=12)
    ax.set_title(title, fontsize=14, fontweight='bold')
    
    # 图例放在右侧
    ax.legend(
        loc='center left',
        bbox_to_anchor=(1.02, 0.5),
        fontsize=10,
        frameon=True,
    )
    
    # 网格
    ax.grid(True, alpha=0.3)
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图片
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"Plot saved to: {output_path}")
    
    plt.close()


def find_best_checkpoint(df: pd.DataFrame, benchmarks: List[str]) -> None:
    """
    找到平均分数最高的 checkpoint 并打印所有 benchmark 结果。
    """
    if df.empty:
        return
    
    # 过滤存在的 benchmarks
    available_benchmarks = [b for b in benchmarks if b in df.columns]
    if not available_benchmarks:
        return
    
    # 计算每个 checkpoint 的平均分数
    df_copy = df.copy()
    df_copy['avg_score'] = df_copy[available_benchmarks].mean(axis=1)
    
    # 只考虑没有缺失 benchmark 的 checkpoint (所有 benchmark 都有值)
    complete_mask = df_copy[available_benchmarks].notna().all(axis=1)
    df_complete = df_copy[complete_mask]
    
    if df_complete.empty:
        print("\n" + "=" * 60)
        print("⚠️  No checkpoint has complete results for all benchmarks")
        print("=" * 60)
        return
    
    # 找到最佳 checkpoint (只从有完整结果的 checkpoint 中选择)
    best_idx = df_complete['avg_score'].idxmax()
    best_row = df_complete.loc[best_idx]
    
    print("\n" + "=" * 60)
    print(f"🏆 Best Checkpoint: checkpoint-{int(best_row['step'])}")
    print(f"   Average Score: {best_row['avg_score']:.2f}%")
    print("=" * 60)
    print("\nBenchmark Results:")
    print("-" * 40)
    
    # 按分数排序打印各 benchmark 结果
    scores = [(b, best_row[b]) for b in available_benchmarks if pd.notna(best_row[b])]
    scores.sort(key=lambda x: x[1], reverse=True)
    
    for benchmark, score in scores:
        print(f"  {benchmark:<20}: {score:.1f}%")
    
    print("-" * 40)
    print(f"  {'Average':<20}: {best_row['avg_score']:.1f}%")
    print("=" * 60)


def save_summary_csv(df: pd.DataFrame, output_path: Path):
    """保存汇总结果为 CSV。"""
    df.to_csv(output_path, index=False)
    print(f"Summary CSV saved to: {output_path}")


def save_tokens_per_step_json(tps_per_checkpoint: List[dict], output_path: Path):
    """保存每个 checkpoint 所有任务的 tokens_per_step 平均信息到 JSON。"""
    if not tps_per_checkpoint:
        return
    out = {
        "description": "每个 checkpoint 下所有任务的 tokens_per_step 汇总（平均等信息）",
        "checkpoints": tps_per_checkpoint,
    }
    output_path.write_text(json.dumps(out, indent=2, ensure_ascii=False))
    print(f"Tokens-per-step summary JSON saved to: {output_path}")


def plot_tokens_per_step_curve(
    df: pd.DataFrame,
    output_path: Path,
    title: str = "Avg Tokens Per Step vs Training Steps",
    figsize: Tuple[int, int] = (10, 6),
):
    """绘制各 checkpoint 的 tokens_per_step 平均均值随 step 变化的折线图。"""
    if df.empty or "tokens_per_step_avg_mean" not in df.columns:
        return
    mask = df["tokens_per_step_avg_mean"].notna()
    if not mask.any():
        return
    fig, ax = plt.subplots(figsize=figsize)
    ax.plot(
        df.loc[mask, "step"],
        df.loc[mask, "tokens_per_step_avg_mean"],
        marker="o",
        markersize=5,
        linewidth=2,
        color="#4363d8",
        label="avg mean (over tasks)",
    )
    if "tokens_per_step_avg_std" in df.columns and df.loc[mask, "tokens_per_step_avg_std"].notna().any():
        ax.fill_between(
            df.loc[mask, "step"],
            df.loc[mask, "tokens_per_step_avg_mean"] - df.loc[mask, "tokens_per_step_avg_std"],
            df.loc[mask, "tokens_per_step_avg_mean"] + df.loc[mask, "tokens_per_step_avg_std"],
            alpha=0.2,
            color="#4363d8",
        )
    ax.set_xlabel("Training Steps", fontsize=12)
    ax.set_ylabel("Tokens Per Step (avg over tasks)", fontsize=12)
    ax.set_title(title, fontsize=14, fontweight="bold")
    ax.legend(loc="best", fontsize=10)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches="tight")
    print(f"Tokens-per-step plot saved to: {output_path}")
    plt.close()


def save_summary_markdown(df: pd.DataFrame, benchmarks: List[str], output_path: Path):
    """保存汇总结果为 Markdown 表格，包含 Avg 列。"""
    if df.empty:
        return
    
    # 确保有可用的 benchmarks
    available_benchmarks = [b for b in benchmarks if b in df.columns]
    
    lines = []
    lines.append("# Checkpoint Benchmark Results\n")
    
    # 表头（加上 Avg 列，以及 TPS(avg) 若存在）
    has_tps = "tokens_per_step_avg_mean" in df.columns
    header = "| Step |"
    separator = "|------|"
    for b in available_benchmarks:
        header += f" {b} |"
        separator += "------|"
    header += " Avg |"
    separator += "------|"
    if has_tps:
        header += " TPS(avg) |"
        separator += "--------|"
    lines.append(header)
    lines.append(separator)

    # 数据行
    for _, row in df.iterrows():
        line = f"| {int(row['step'])} |"
        valid_scores = []
        has_missing = False
        for b in available_benchmarks:
            if b in row and pd.notna(row[b]):
                line += f" {row[b]:.1f} |"
                valid_scores.append(row[b])
            else:
                line += " - |"
                has_missing = True

        # 计算 Avg - 只有当所有 benchmark 都有值时才计算
        if valid_scores and not has_missing:
            avg = sum(valid_scores) / len(valid_scores)
            line += f" {avg:.1f} |"
        else:
            line += " - |"
        if has_tps and pd.notna(row.get("tokens_per_step_avg_mean")):
            line += f" {row['tokens_per_step_avg_mean']:.3f} |"
        elif has_tps:
            line += " - |"
        lines.append(line)
    
    lines.append("")
    
    output_path.write_text('\n'.join(lines))
    print(f"Summary markdown saved to: {output_path}")


def main():
    parser = argparse.ArgumentParser(
        description="汇总多个 checkpoint 的评测结果并画图",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python plot_checkpoint_results.py --input_dir VLMEvalKit_local/output_models/xxx
    python plot_checkpoint_results.py --input_dir VLMEvalKit_local/output_models/xxx --benchmarks "GQA,POPE"
        """,
    )
    
    parser.add_argument(
        "--input_dir",
        type=str,
        required=True,
        help="包含 eval_vlmeval/checkpoint-*/summary.md 的目录",
    )
    parser.add_argument(
        "--benchmarks",
        type=str,
        default=None,
        help="逗号分隔的 benchmark 列表 (默认: 所有 benchmarks)",
    )
    parser.add_argument(
        "--output_prefix",
        type=str,
        default="checkpoint_results",
        help="输出文件前缀 (default: checkpoint_results)",
    )
    parser.add_argument(
        "--title",
        type=str,
        default=None,
        help="图表标题 (default: 从目录名推断)",
    )
    
    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
    
    print(f"Collecting results from: {input_dir}")
    
    # 收集所有结果（含每个 checkpoint 的 tokens_per_step 平均信息）
    df, all_benchmarks, tps_per_checkpoint = collect_all_results(input_dir)

    if df.empty:
        print("No results found.")
        return 1

    print(f"\nFound {len(df)} checkpoints with {len(all_benchmarks)} benchmarks")
    if tps_per_checkpoint:
        print(f"Tokens-per-step: {len(tps_per_checkpoint)} checkpoints with TPS stats")
    print(f"Steps range: {df['step'].min()} - {df['step'].max()}")
    print(f"Benchmarks: {', '.join(all_benchmarks)}")
    
    # 确定要绘制的 benchmarks
    if args.benchmarks:
        selected_benchmarks = [b.strip() for b in args.benchmarks.split(',')]
    else:
        selected_benchmarks = all_benchmarks
    
    # 输出目录与输入目录相同
    output_dir = input_dir
    
    # 生成标题
    if args.title:
        title = args.title
    else:
        title = f"Benchmark Scores: {input_dir.name}"
    
    # 保存结果
    output_prefix = args.output_prefix
    
    # 1. 保存 CSV
    csv_path = output_dir / f"{output_prefix}.csv"
    save_summary_csv(df, csv_path)
    
    # 2. 保存 Markdown
    md_path = output_dir / f"{output_prefix}.md"
    save_summary_markdown(df, selected_benchmarks, md_path)
    
    # 3. 绘制并保存图表
    plot_path = output_dir / f"{output_prefix}.png"
    plot_benchmark_curves(
        df,
        selected_benchmarks,
        plot_path,
        title=title,
    )

    # 4. 保存每个 checkpoint 所有任务的 tokens_per_step 平均信息 JSON，并绘制 TPS 曲线
    tps_json_path = output_dir / f"{output_prefix}_tokens_per_step.json"
    save_tokens_per_step_json(tps_per_checkpoint, tps_json_path)
    tps_plot_path = output_dir / f"{output_prefix}_tokens_per_step.png"
    plot_tokens_per_step_curve(
        df,
        tps_plot_path,
        title=f"Avg Tokens Per Step vs Steps: {input_dir.name}",
    )

    # 5. 打印最佳 checkpoint
    find_best_checkpoint(df, selected_benchmarks)
    
    print("\nDone!")
    return 0


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

