#!/usr/bin/env python3
"""
Collect benchmark results from an EvalScope output directory and write a summary
Excel/CSV similar to:

  P800模型能力评测结果 - DS4-Flash-INT8-NO-Thinking-2.0-FULL.csv

Rules:
- Scores are averaged across seeds / multi-runs.
- Perf metrics (latency, TTFT, TPOT, TPS, tokens) are recomputed from raw
  predictions across all seeds / multi-runs, so resuming from a checkpoint
  no longer resets cumulative statistics.
"""

import argparse
import json
import statistics
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd


# Map benchmark -> capability domain (same as the reference CSV)
BENCHMARK_DOMAIN = {
    'bigcodebench': '代码与工程',
    'humaneval': '代码与工程',
    'live_code_bench': '代码与工程',
    'aime24': '推理与数学',
    'aime25': '推理与数学',
    'aime26': '推理与数学',
    'hmmt26': '推理与数学',
    'imo_answerbench': '推理与数学',
    'gsm8k': '推理与数学',
    'competition_math': '推理与数学',
    'bbh': '推理与数学',
    'drop': '推理与数学',
    'gpqa_diamond': '知识与语言理解',
    'hle': '知识与语言理解',
    'hle': '知识与语言理解',
    'hle_low': '知识与语言理解',
    'mmlu_pro': '知识与语言理解',
    'simple_qa': '知识与语言理解',
    'super_gpqa': '知识与语言理解',
    'mmlu': '知识与语言理解',
    'cmmlu': '知识与语言理解',
    'arc': '知识与语言理解',
    'hellaswag': '知识与语言理解',
    'trivia_qa': '知识与语言理解',
    'winogrande': '知识与语言理解',
    'longbench_v2': '长上下文',
    'openai_mrcr': '长上下文',
    'tau2_bench': '智能体与工具',
    'general_fc': '智能体与工具',
    'bfcl_v3': '智能体与工具',
    'terminal_bench_v2_1': '智能体与工具',
    # 指纹/安全类 benchmark（bash/fingerprint/ 下的独立执行器产出）
    'llmmap': '模型安全与指纹',
    'llm_verify': '模型安全与指纹',
    'llm_fingerprint_detector': '模型安全与指纹',
    'fp_fusion': '模型安全与指纹',
}

# Column order matching the reference CSV
OUTPUT_COLUMNS = [
    '分类',
    'Benchmark',
    '得分',
    '实测时间(h)',
    '总样本数',
    '延迟_mean(s)',
    '输出TPS',
    '请求QPS',
    '输入tokens_mean',
    '输出tokens_mean',
    '累计总tokens',
    'TTFT_mean(s)',
    'TTFT P90',
    'TTFT P99',
    'TPOT_mean(s)',
    'TPOT P90',
    'TPOT P99',
]

# Some evalscope output directory names do not match the canonical benchmark
# name (e.g. the hle benchmark writes to the `hle_low` directory because of
# subset naming). Map them to the canonical name here so the summary uses a
# consistent label.
BENCHMARK_NAME_ALIAS = {
    'hle_low': 'hle',
}


def percentile(values, q):
    """Return the q-th percentile using numpy's linear interpolation."""
    if not values:
        return np.nan
    return float(np.percentile(values, q))


def read_predictions(pred_file: Path):
    """Yield perf_metrics dicts from a predictions JSONL file."""
    for obj in read_predictions_with_index(pred_file):
        if obj['perf_metrics'] is not None:
            yield obj['perf_metrics']


def read_predictions_with_index(pred_file: Path):
    """Yield ``{'index', 'perf_metrics'}`` dicts from a predictions JSONL file."""
    if not pred_file.exists():
        return
    with open(pred_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue
            model_output = obj.get('model_output', {})
            pm = model_output.get('perf_metrics')
            if not pm and 'choices' in model_output:
                choices = model_output['choices']
                if choices and 'message' in choices[0]:
                    pm = choices[0]['message'].get('perf_metrics')
            yield {'index': obj.get('index'), 'perf_metrics': pm}


def read_agent_perf_from_trajectory(pred_file: Path):
    """Extract approximate perf metrics from agent trajectory files.

    Agent/sandbox benchmarks (e.g. terminal_bench_v2_1) do not record per-call
    latency/TTFT/TPOT through EvalScope's model wrapper. However, the trial
    trajectory contains step timestamps and final token counts. This function
    yields synthetic ``{'index', 'perf_metrics'}`` records with:

      - ``latency``: wall-clock trial duration in seconds
      - ``input_tokens``: total prompt tokens from the agent run
      - ``output_tokens``: total completion tokens from the agent run
      - ``ttft`` / ``tpot``: not available, left as None
    """
    if not pred_file.exists():
        return
    from datetime import datetime
    with open(pred_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue
            idx = obj.get('index')
            model_output = obj.get('model_output', {})
            content = ''
            choices = model_output.get('choices', [])
            if choices and 'message' in choices[0]:
                content = choices[0]['message'].get('content', '')
            if not content or not isinstance(content, str) or not content.startswith('file://'):
                continue
            traj_path = Path(content[7:]) / 'agent' / 'trajectory.json'
            if not traj_path.exists():
                continue
            try:
                traj = json.loads(traj_path.read_text(encoding='utf-8'))
            except Exception:
                continue
            steps = traj.get('steps', [])
            duration = None
            if len(steps) >= 2:
                try:
                    first_ts = datetime.fromisoformat(steps[0]['timestamp'])
                    last_ts = datetime.fromisoformat(steps[-1]['timestamp'])
                    duration = (last_ts - first_ts).total_seconds()
                except Exception:
                    pass
            final_metrics = traj.get('final_metrics', {})
            prompt_tokens = final_metrics.get('total_prompt_tokens')
            completion_tokens = final_metrics.get('total_completion_tokens')
            pm = {}
            if duration is not None:
                pm['latency'] = duration
            if prompt_tokens is not None:
                pm['input_tokens'] = int(prompt_tokens)
            if completion_tokens is not None:
                pm['output_tokens'] = int(completion_tokens)
            if pm:
                yield {'index': idx, 'perf_metrics': pm}


def load_backup_summary(output_dir: Path, benchmark: str, model_name: str):
    """Load the durable ``perf_stats_backup/<benchmark>__<model>.json``.

    Returns the ``summary`` dict (with ``latency`` / ``ttft`` / ``tpot`` /
    ``throughput`` / ``usage`` sub-dicts) or ``None`` if no backup exists.
    """
    try:
        from perf_backup import get_backup_paths
    except ImportError:
        return None
    perf_path, _ = get_backup_paths(Path(output_dir), benchmark, model_name)
    if not perf_path.exists():
        return None
    try:
        payload = json.loads(perf_path.read_text(encoding='utf-8'))
        return payload.get('summary')
    except Exception:
        return None


def find_report(output_dir: Path, benchmark: str, model_name: str):
    """Find the first report JSON for a benchmark/model under output_dir."""
    bench_dir = output_dir / benchmark
    if not bench_dir.exists():
        return None
    for seed_dir in sorted(bench_dir.iterdir()):
        if not seed_dir.is_dir():
            continue
        for candidate in find_report_candidates(seed_dir, benchmark, model_name):
            if candidate.exists():
                return candidate
    return None


def find_report_candidates(seed_dir: Path, benchmark: str, model_name: str):
    """Return candidate report paths, accepting variant naming patterns."""
    # The actual report filename may not match the directory name exactly
    # (e.g. hle_low directory holds hle.json). Try all .json files under
    # the reports directory.
    candidates = [
        seed_dir / 'reports' / f'{benchmark}.json',
    ]
    for p in (seed_dir / 'reports').glob('*.json'):
        candidates.append(p)
    return candidates


def find_all_reports(output_dir: Path, benchmark: str, model_name: str):
    """Find all report JSONs for a benchmark/model (across seeds/runs)."""
    bench_dir = output_dir / benchmark
    reports = []
    if not bench_dir.exists():
        return reports
    for seed_dir in sorted(bench_dir.iterdir()):
        if not seed_dir.is_dir():
            continue
        for candidate in find_report_candidates(seed_dir, benchmark, model_name):
            if candidate.exists() and candidate not in reports:
                reports.append(candidate)
                break
    return reports


def find_archive_predictions(output_dir: Path, benchmark: str, model_name: str):
    """Return the durable per-sample archive maintained by ``perf_backup.py``.

    The archive contains every sample ever produced across all runs and
    breakpoints for this benchmark/model (deduplicated by ``index``). When
    present, the summary aggregator should prefer it over the raw
    ``predictions/*.jsonl`` files because the latter can be overwritten on
    restart.
    """
    try:
        from perf_backup import get_backup_paths
    except ImportError:
        return []
    _, archive_path = get_backup_paths(output_dir, benchmark, model_name)
    if archive_path.exists():
        return [archive_path]
    return []


def find_all_predictions(output_dir: Path, benchmark: str, model_name: str):
    """Find all predictions JSONL files for a benchmark/model.

    For single-seed runs the durable ``predictions_archive`` is preferred so
    breakpoint-resume does not lose completed samples.  When multiple seed/run
    directories exist (multi-run benchmarks) we aggregate from each run
    separately and skip the archive, because the archive only keeps the latest
    record per ``index`` and would otherwise collide with one of the runs.
    """
    bench_dir = output_dir / benchmark
    seed_dirs = []
    if bench_dir.exists():
        seed_dirs = sorted([p for p in bench_dir.iterdir() if p.is_dir()])

    files = []
    # Only rely on the archive for single-run / resume scenarios.
    if len(seed_dirs) <= 1:
        files = list(find_archive_predictions(output_dir, benchmark, model_name))

    for seed_dir in seed_dirs:
        pred_dir = seed_dir / 'predictions'
        if pred_dir.exists():
            files.extend(sorted(pred_dir.rglob('*.jsonl')))
    return files


def parse_log_duration(log_file: Path):
    """Parse first and last timestamp from eval_log.log and return duration in hours."""
    if not log_file.exists():
        return np.nan
    from datetime import datetime
    first_dt = None
    last_dt = None
    fmt = '%Y-%m-%d %H:%M:%S'
    with open(log_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if len(line) < 19:
                continue
            try:
                dt = datetime.strptime(line[:19], fmt)
            except ValueError:
                continue
            if first_dt is None:
                first_dt = dt
            last_dt = dt
    if first_dt is None or last_dt is None or last_dt <= first_dt:
        return np.nan
    return (last_dt - first_dt).total_seconds() / 3600.0


def _metric_name(metric: dict) -> Optional[str]:
    """Return a metric's display/key name across report schema v1 and v2."""
    if not isinstance(metric, dict):
        return None
    if metric.get('name'):
        return str(metric['name'])
    identity = metric.get('identity') or {}
    if isinstance(identity, dict) and identity.get('name'):
        return str(identity['name'])
    if metric.get('legacy_name'):
        return str(metric['legacy_name'])
    return None


def _identity_key(identity: Optional[dict]) -> Optional[tuple]:
    if not isinstance(identity, dict) or not identity.get('name'):
        return None
    dims = identity.get('dimensions') or {}
    if not isinstance(dims, dict):
        dims = {}
    return (
        str(identity.get('name')),
        str(identity.get('aggregation') or 'mean'),
        tuple(sorted((str(k), str(v)) for k, v in dims.items())),
    )


def extract_score(report_data: dict) -> float:
    """Extract the primary score from a report JSON.

    Supports:
    - legacy reports with top-level ``score`` / metrics named ``mean_acc``
    - EvalScope report schema v2 with ``primary_metric_identity`` +
      ``metrics[].identity`` / ``metrics[].score``
    """
    score = report_data.get('score')
    if score is not None:
        return float(score)

    metrics = report_data.get('metrics') or []
    if not metrics:
        return 0.0

    # Schema v2: prefer the explicit primary metric identity when present.
    primary_identity = report_data.get('primary_metric_identity')
    primary_key = _identity_key(primary_identity)
    if primary_key is not None:
        for m in metrics:
            if _identity_key(m.get('identity')) == primary_key:
                return float(m.get('score', m.get('macro_score', 0.0)))

    # Legacy / fallback preferred names.
    preferred = {
        'mean_acc',
        'accuracy',
        'acc',
        'main_problem_pass_rate',
        'pass_rate',
        'normalized_score',
        'f1',
    }
    for m in metrics:
        name = _metric_name(m)
        if name in preferred:
            return float(m.get('score', m.get('macro_score', 0.0)))

    # Last resort: first metric with a numeric score.
    for m in metrics:
        if m.get('score') is not None:
            return float(m.get('score'))
        if m.get('macro_score') is not None:
            return float(m.get('macro_score'))
    return 0.0


def collect_benchmark(output_dir: Path, benchmark: str, model_name: str):
    """Collect aggregated results for one benchmark."""
    reports = find_all_reports(output_dir, benchmark, model_name)
    if not reports:
        return None

    scores = []
    summary0 = None
    n_samples_unique = 0
    req_success = 0
    req_failed = 0
    req_client = 0
    for report in reports:
        try:
            data = json.loads(report.read_text(encoding='utf-8'))
            scores.append(extract_score(data))
            if summary0 is None:
                perf_metrics = data.get('perf_metrics') or {}
                summary0 = perf_metrics.get('summary', {})
                n_samples_unique = summary0.get('n_samples', data.get('num', 0))
            req = ((data.get('perf_metrics') or {}).get('summary') or {}).get('request') or {}
            req_success += int(req.get('success_attempts') or 0)
            req_failed += int(req.get('failed_attempts') or 0)
            req_client += int(req.get('client_errors') or 0)
        except Exception:
            continue

    avg_score = float(np.mean(scores)) if scores else 0.0

    # Aggregate raw prediction perf metrics across all seeds/runs.
    # This fixes the breakpoint-resume issue where cumulative stats are reset.
    # Deduplicate by (run, sample `index`) so that:
    #   1. multiple prediction jsonl files inside the same run directory
    #      (e.g. `<benchmark>__<model>.jsonl` and `<benchmark>_<subset>.jsonl`)
    #      do not double-count the same sample;
    #   2. each seed/run still contributes its own predictions for multi-run
    #      benchmarks, so total sample count is sum(runs).
    pred_files = find_all_predictions(output_dir, benchmark, model_name)
    seen_keys = set()
    latencies = []
    ttfts = []
    tpots = []
    input_tokens = []
    output_tokens = []
    sample_indexes = []
    for pf in pred_files:
        # ``run_key`` is the seed/run directory name, or the archive filename
        # for archive-only scenarios.
        run_key = pf.name
        parts = pf.parts
        if 'predictions_archive' not in parts:
            for i, part in enumerate(parts):
                if part == 'predictions' and i > 0:
                    run_key = parts[i - 1]
                    break
        for obj in read_predictions_with_index(pf):
            idx = obj['index']
            key = (run_key, idx)
            if idx is None:
                # Keep perf data even when we lack an index, so older
                # benchmark files without `index` don't get dropped.
                pm = obj['perf_metrics']
            elif key in seen_keys:
                continue
            else:
                seen_keys.add(key)
                sample_indexes.append(idx)
                pm = obj['perf_metrics']
            if pm is None:
                continue
            if pm.get('latency') is not None:
                latencies.append(float(pm['latency']))
            if pm.get('ttft') is not None:
                ttfts.append(float(pm['ttft']))
            if pm.get('tpot') is not None:
                tpots.append(float(pm['tpot']))

            itok = pm.get('input_tokens')
            otok = pm.get('output_tokens')
            if (itok is None or otok is None) and 'usage' in pm:
                itok = pm['usage'].get('input_tokens') if itok is None else itok
                otok = pm['usage'].get('output_tokens') if otok is None else otok
            if itok is not None:
                input_tokens.append(int(itok))
            if otok is not None:
                output_tokens.append(int(otok))

    # Agent/sandbox benchmarks do not expose per-call perf metrics through
    # EvalScope's model wrapper. Fall back to parsing the trial trajectory
    # files (timestamps + final token counts) to get approximate latency and
    # token usage.
    if not latencies and pred_files:
        for obj in read_agent_perf_from_trajectory(pred_files[0]):
            idx = obj['index']
            key = ('trajectory', idx)
            if idx is not None:
                if key in seen_keys:
                    continue
                seen_keys.add(key)
                sample_indexes.append(idx)
            pm = obj['perf_metrics']
            if pm.get('latency') is not None:
                latencies.append(float(pm['latency']))
            if pm.get('input_tokens') is not None:
                input_tokens.append(int(pm['input_tokens']))
            if pm.get('output_tokens') is not None:
                output_tokens.append(int(pm['output_tokens']))

    # If we don't have raw predictions but have a perf_stats backup, that
    # represents a known-good summary captured right after a clean run —
    # preferable to summary0 (which may be the just-reset run).
    backup_summary = None
    if not latencies:
        backup_summary = load_backup_summary(output_dir, benchmark, model_name)

    if latencies:
        latency_mean = float(np.mean(latencies))
        total_compute_time = float(np.sum(latencies))
        total_output_tokens = sum(output_tokens)
        avg_output_tps = total_output_tokens / total_compute_time if total_compute_time > 0 else np.nan
        avg_req_ps = len(latencies) / total_compute_time if total_compute_time > 0 else np.nan
        input_tok_mean = float(np.mean(input_tokens)) if input_tokens else np.nan
        output_tok_mean = float(np.mean(output_tokens)) if output_tokens else np.nan
        total_tokens = sum(input_tokens) + sum(output_tokens)
        ttft_mean = float(np.mean(ttfts)) if ttfts else np.nan
        ttft_p90 = percentile(ttfts, 90)
        ttft_p99 = percentile(ttfts, 99)
        tpot_mean = float(np.mean(tpots)) if tpots else np.nan
        tpot_p90 = percentile(tpots, 90)
        tpot_p99 = percentile(tpots, 99)
        # The actual sample count we just rebuilt from raw predictions is more
        # reliable than the (possibly reset) report summary's n_samples.
        # For multi-seed / multi-run benchmarks we report the total number of
        # evaluated predictions (all seeds/runs) rather than unique problem IDs.
        if n_samples_unique < len(latencies):
            n_samples_unique = len(latencies)
    elif summary0:
        # Fallback to report summary if raw predictions are unavailable
        latency_mean = summary0.get('latency', {}).get('mean', np.nan)
        avg_output_tps = summary0.get('throughput', {}).get('avg_output_tps', np.nan)
        avg_req_ps = summary0.get('throughput', {}).get('avg_req_ps', np.nan)
        input_tok_mean = summary0.get('tokens', {}).get('input_tokens', {}).get('mean', np.nan)
        output_tok_mean = summary0.get('tokens', {}).get('output_tokens', {}).get('mean', np.nan)
        total_tokens = summary0.get('tokens', {}).get('total_tokens_count', np.nan)
        ttft_mean = summary0.get('ttft', {}).get('mean', np.nan)
        ttft_p90 = summary0.get('ttft', {}).get('90%', np.nan)
        ttft_p99 = summary0.get('ttft', {}).get('99%', np.nan)
        tpot_mean = summary0.get('tpot', {}).get('mean', np.nan)
        tpot_p90 = summary0.get('tpot', {}).get('90%', np.nan)
        tpot_p99 = summary0.get('tpot', {}).get('99%', np.nan)
        # If the just-read report looks like a freshly-reset run (smaller
        # n_samples than the durable backup), prefer the backup's summary so
        # the cumulative numbers are not lost.
        try:
            from perf_backup import get_backup_paths
            backup_path, _ = get_backup_paths(output_dir, benchmark, model_name)
            if backup_path.exists():
                payload = json.loads(backup_path.read_text(encoding='utf-8'))
                if (payload.get('n_samples') or 0) > (summary0.get('n_samples') or 0):
                    backup_summary = payload.get('summary') or {}
                    latency_mean = backup_summary.get('latency', {}).get('mean', latency_mean)
                    avg_output_tps = backup_summary.get('throughput', {}).get('avg_output_tps', avg_output_tps)
                    avg_req_ps = backup_summary.get('throughput', {}).get('avg_req_ps', avg_req_ps)
                    input_tok_mean = backup_summary.get('usage', {}).get('input_tokens', {}).get('mean', input_tok_mean)
                    output_tok_mean = backup_summary.get('usage', {}).get('output_tokens', {}).get('mean', output_tok_mean)
                    total_tokens = backup_summary.get('usage', {}).get('total_tokens_count', total_tokens)
                    ttft_mean = backup_summary.get('ttft', {}).get('mean', ttft_mean)
                    ttft_p90 = backup_summary.get('ttft', {}).get('90%', ttft_p90)
                    ttft_p99 = backup_summary.get('ttft', {}).get('99%', ttft_p99)
                    tpot_mean = backup_summary.get('tpot', {}).get('mean', tpot_mean)
                    tpot_p90 = backup_summary.get('tpot', {}).get('90%', tpot_p90)
                    tpot_p99 = backup_summary.get('tpot', {}).get('99%', tpot_p99)
                    n_samples_unique = payload.get('n_samples') or n_samples_unique
        except Exception:
            pass
    elif backup_summary:
        latency_mean = backup_summary.get('latency', {}).get('mean', np.nan)
        avg_output_tps = backup_summary.get('throughput', {}).get('avg_output_tps', np.nan)
        avg_req_ps = backup_summary.get('throughput', {}).get('avg_req_ps', np.nan)
        input_tok_mean = backup_summary.get('usage', {}).get('input_tokens', {}).get('mean', np.nan)
        output_tok_mean = backup_summary.get('usage', {}).get('output_tokens', {}).get('mean', np.nan)
        total_tokens = backup_summary.get('usage', {}).get('total_tokens_count', np.nan)
        ttft_mean = backup_summary.get('ttft', {}).get('mean', np.nan)
        ttft_p90 = backup_summary.get('ttft', {}).get('90%', np.nan)
        ttft_p99 = backup_summary.get('ttft', {}).get('99%', np.nan)
        tpot_mean = backup_summary.get('tpot', {}).get('mean', np.nan)
        tpot_p90 = backup_summary.get('tpot', {}).get('90%', np.nan)
        tpot_p99 = backup_summary.get('tpot', {}).get('99%', np.nan)
    else:
        # Agent / sandbox benchmarks (e.g. terminal_bench_v2_1) provide a score
        # but do not record per-sample perf metrics. Keep the score and leave
        # perf columns empty rather than dropping the benchmark entirely.
        latency_mean = np.nan
        avg_output_tps = np.nan
        avg_req_ps = np.nan
        input_tok_mean = np.nan
        output_tok_mean = np.nan
        total_tokens = np.nan
        ttft_mean = np.nan
        ttft_p90 = np.nan
        ttft_p99 = np.nan
        tpot_mean = np.nan
        tpot_p90 = np.nan
        tpot_p99 = np.nan
        if not n_samples_unique and reports:
            try:
                data = json.loads(reports[0].read_text(encoding='utf-8'))
                n_samples_unique = data.get('num', 0)
            except Exception:
                pass

    # Duration: prefer the durable active timer maintained by perf_backup.py,
    # which only counts time when run_task() is actually executing. This avoids
    # counting idle gaps caused by manual interruption/resume.
    # If no active timer exists, fall back to log duration or latency estimate.
    active_time_hours = np.nan
    try:
        from perf_backup import load_active_time
        active_seconds = load_active_time(output_dir, benchmark, model_name)
        if active_seconds > 0:
            active_time_hours = active_seconds / 3600.0
    except Exception:
        pass

    if not np.isnan(active_time_hours):
        duration_hours = active_time_hours
    else:
        # Legacy fallback: sum of wall-clock durations from eval_log.log.
        # For sandbox benchmarks the main log may not cover sandbox execution,
        # so we also estimate wall time as compute_time/batch_size and take max.
        duration_hours = 0.0
        batch_size = 4
        for report in reports:
            log_file = report.parent.parent.parent / 'logs' / 'eval_log.log'
            d = parse_log_duration(log_file)
            if not np.isnan(d):
                duration_hours += d
            cfg_file = report.parent.parent.parent / 'configs' / 'task_config.yaml'
            if cfg_file.exists():
                try:
                    import yaml
                    cfg = yaml.safe_load(cfg_file.read_text(encoding='utf-8'))
                    batch_size = int(cfg.get('eval_batch_size', batch_size))
                except Exception:
                    pass

        compute_wall_estimate = (total_compute_time / batch_size / 3600.0) if latencies and batch_size > 0 else np.nan
        if duration_hours <= 0:
            duration_hours = compute_wall_estimate
        else:
            duration_hours = max(duration_hours, compute_wall_estimate)

    vendor_http = req_success + req_failed
    request_success_rate = (req_success / vendor_http) if vendor_http > 0 else np.nan

    return {
        '分类': BENCHMARK_DOMAIN.get(benchmark, '其他'),
        'Benchmark': BENCHMARK_NAME_ALIAS.get(benchmark, benchmark),
        '得分': round(avg_score, 4),
        '实测时间(h)': round(duration_hours, 4) if not np.isnan(duration_hours) else np.nan,
        '总样本数': n_samples_unique,
        '请求成功率': round(request_success_rate, 4) if not np.isnan(request_success_rate) else np.nan,
        'HTTP成功': req_success if vendor_http or req_client else np.nan,
        'HTTP失败': req_failed if vendor_http or req_client else np.nan,
        'client_errors': req_client if vendor_http or req_client else np.nan,
        '延迟_mean(s)': round(latency_mean, 5) if not np.isnan(latency_mean) else np.nan,
        '输出TPS': round(avg_output_tps, 2) if not np.isnan(avg_output_tps) else np.nan,
        '请求QPS': round(avg_req_ps, 4) if not np.isnan(avg_req_ps) else np.nan,
        '输入tokens_mean': round(input_tok_mean, 2) if not np.isnan(input_tok_mean) else np.nan,
        '输出tokens_mean': round(output_tok_mean, 2) if not np.isnan(output_tok_mean) else np.nan,
        '累计总tokens': total_tokens if not np.isnan(total_tokens) else np.nan,
        'TTFT_mean(s)': round(ttft_mean, 5) if not np.isnan(ttft_mean) else np.nan,
        'TTFT P90': round(ttft_p90, 5) if not np.isnan(ttft_p90) else np.nan,
        'TTFT P99': round(ttft_p99, 5) if not np.isnan(ttft_p99) else np.nan,
        'TPOT_mean(s)': round(tpot_mean, 5) if not np.isnan(tpot_mean) else np.nan,
        'TPOT P90': round(tpot_p90, 5) if not np.isnan(tpot_p90) else np.nan,
        'TPOT P99': round(tpot_p99, 5) if not np.isnan(tpot_p99) else np.nan,
    }


def eval_benchmark(benchmark_names: list, output_dir: Path, model_name: str,
                   out_name: str = None, excel_output_dir: Path = None):
    """Collect results for a specific list of benchmarks.

    This is the public entry point intended for use by ``run.py`` and other
    scripts that already know which benchmarks they ran.

    Args:
        benchmark_names: List of canonical benchmark names to aggregate (e.g.
            ``['aime24', 'gsm8k', 'arc']``). Only benchmarks with reports on
            disk will appear in the summary.
        output_dir: EvalScope output root directory.
        model_name: Model name to look up reports/predictions for.
        out_name: Output file name (without extension). Defaults to the safe
            model name.
        excel_output_dir: Optional separate directory for the Excel summary.
            If provided, the Excel file is written here while the CSV stays
            under ``output_dir/results/<model>/``.

    Returns:
        (csv_path, xlsx_path) tuple.
    """
    return collect_all(output_dir, model_name, out_name,
                       include_benchmarks=benchmark_names,
                       excel_output_dir=excel_output_dir)


def collect_all(output_dir: Path, model_name: str, out_name: str = None,
                include_benchmarks: list = None,
                excel_output_dir: Path = None):
    """Collect benchmarks under ``output_dir`` and write summary Excel/CSV."""
    if not output_dir.exists():
        raise FileNotFoundError(f'Output directory not found: {output_dir}')

    safe_model = model_name.replace('/', '_').replace('\\', '_').replace(' ', '_')

    # Discover benchmarks
    if include_benchmarks is not None:
        # Whitelist mode: only aggregate benchmarks that were actually run in
        # this invocation. Directory names may differ from canonical names
        # (e.g. evalscope writes the hle benchmark to the hle_low directory),
        # so we canonicalize each directory name via BENCHMARK_NAME_ALIAS and
        # check membership in the caller's list. Order follows the caller's
        # include_benchmarks list.
        include_set = set(include_benchmarks)
        discovered = []  # list of (dir_name, canonical_name)
        for bench_dir in output_dir.iterdir():
            if not bench_dir.is_dir() or bench_dir.name == safe_model:
                continue
            if not find_all_reports(output_dir, bench_dir.name, model_name):
                continue
            canonical = BENCHMARK_NAME_ALIAS.get(bench_dir.name, bench_dir.name)
            if canonical in include_set or bench_dir.name in include_set:
                discovered.append((bench_dir.name, canonical))

        order = {name: idx for idx, name in enumerate(include_benchmarks)}
        discovered.sort(key=lambda x: order.get(x[1], len(include_benchmarks)))

        benchmarks = [dir_name for dir_name, _ in discovered]
    else:
        # Legacy behaviour: scan everything on disk.
        benchmarks = []
        for bench_dir in sorted(output_dir.iterdir()):
            if not bench_dir.is_dir() or bench_dir.name == safe_model:
                continue
            reports = find_all_reports(output_dir, bench_dir.name, model_name)
            if reports:
                benchmarks.append(bench_dir.name)

    rows = []
    for benchmark in benchmarks:
        row = collect_benchmark(output_dir, benchmark, model_name)
        if row:
            rows.append(row)

    if not rows:
        print(f'No results found for model {model_name} under {output_dir}')
        return None, None

    df = pd.DataFrame(rows, columns=OUTPUT_COLUMNS)
    df = df.drop_duplicates(subset=['Benchmark'], keep='first')

    # Add total row
    total_score = df['得分'].mean()
    total_time = df['实测时间(h)'].sum()
    total_samples = df['总样本数'].sum() if '总样本数' in df.columns else np.nan
    total_tokens = df['累计总tokens'].sum() if '累计总tokens' in df.columns else np.nan
    total_row = {
        '分类': '总计',
        'Benchmark': '',
        '得分': round(total_score, 4),
        '实测时间(h)': round(total_time, 4),
        '总样本数': total_samples if not np.isnan(total_samples) else np.nan,
        '累计总tokens': total_tokens if not np.isnan(total_tokens) else np.nan,
    }
    for col in OUTPUT_COLUMNS:
        if col not in total_row:
            total_row[col] = np.nan
    df = pd.concat([df, pd.DataFrame([total_row], columns=OUTPUT_COLUMNS)], ignore_index=True)

    if out_name is None:
        out_name = safe_model

    # Excel can optionally be written to a separate project-level results dir
    # so that the latest summary per model is easy to find independently of
    # the per-run output directory. When a separate dir is given, put both
    # CSV and Excel there.
    if excel_output_dir is not None:
        summary_dir = Path(excel_output_dir)
    else:
        summary_dir = output_dir / 'results' / safe_model
    summary_dir.mkdir(parents=True, exist_ok=True)

    csv_path = summary_dir / f'{out_name}.csv'
    xlsx_path = summary_dir / f'{out_name}.xlsx'

    df.to_csv(csv_path, index=False, encoding='utf-8-sig')
    df.to_excel(xlsx_path, index=False)

    print(f'Summary written to:')
    print(f'  CSV:  {csv_path}')
    print(f'  Excel: {xlsx_path}')
    return csv_path, xlsx_path


def main():
    parser = argparse.ArgumentParser(description='Collect EvalScope benchmark results into Excel/CSV')
    parser.add_argument('--output-dir', required=True, help='EvalScope output directory')
    parser.add_argument('--model', default='DeepSeek-V4-Flash-Int8', help='Model name used in reports')
    parser.add_argument('--out-name', default=None, help='Output file name (without extension); defaults to safe model name')
    parser.add_argument('--include', default=None,
                        help='Comma-separated benchmark whitelist (default: all on disk)')
    parser.add_argument('--benchmarks', '--benchmarks', dest='benchmarks', default=None,
                        help='Alias for --include')
    parser.add_argument('--excel-output-dir', default=None,
                        help='Separate directory for the Excel summary; defaults to <project_root>/results')
    args = parser.parse_args()

    # Default Excel output dir to project-level results/ so the latest summary
    # per model is easy to find independently of the per-run output directory.
    if args.excel_output_dir is None:
        excel_output_dir = Path(__file__).parent.parent / 'results'
    else:
        excel_output_dir = Path(args.excel_output_dir)

    include_benchmarks = None
    if args.include:
        include_benchmarks = [b.strip() for b in args.include.split(',') if b.strip()]
    elif args.benchmarks:
        include_benchmarks = [b.strip() for b in args.benchmarks.split(',') if b.strip()]

    if include_benchmarks is not None:
        eval_benchmark(include_benchmarks, Path(args.output_dir), args.model, args.out_name,
                        excel_output_dir=excel_output_dir)
    else:
        collect_all(Path(args.output_dir), args.model, args.out_name,
                    excel_output_dir=excel_output_dir)


if __name__ == '__main__':
    main()