#!/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.
- Writing CSV/Excel upserts by Benchmark name: existing rows for other
  benchmarks are kept, the finished benchmark overwrites its own row, and
  a new benchmark is appended. The 总计 row is always recomputed.
"""

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': '模型安全与指纹',
}

TOTAL_CATEGORY = '总计'

# 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 _usage_block(summary: dict) -> dict:
    """Prefer ``usage`` (current reports) and fall back to legacy ``tokens``."""
    if not isinstance(summary, dict):
        return {}
    usage = summary.get('usage')
    if isinstance(usage, dict) and usage:
        return usage
    tokens = summary.get('tokens')
    return tokens if isinstance(tokens, dict) else {}


def _as_num(value):
    return np.nan if value is None else value


def request_perf_from_summary(summary: Optional[dict]) -> Optional[dict]:
    """Map per-request ``perf_metrics.summary`` onto CSV column values.

    Agent/sandbox reports record TTFT/TPOT/latency per LLM request here, even
    when prediction jsonl rows have no ``perf_metrics``.
    """
    if not isinstance(summary, dict) or not summary:
        return None
    latency = summary.get('latency') if isinstance(summary.get('latency'), dict) else {}
    throughput = summary.get('throughput') if isinstance(summary.get('throughput'), dict) else {}
    ttft = summary.get('ttft') if isinstance(summary.get('ttft'), dict) else {}
    tpot = summary.get('tpot') if isinstance(summary.get('tpot'), dict) else {}
    usage = _usage_block(summary)
    if not any((latency, ttft, tpot, usage, throughput)):
        return None
    in_tok = usage.get('input_tokens') if isinstance(usage.get('input_tokens'), dict) else {}
    out_tok = usage.get('output_tokens') if isinstance(usage.get('output_tokens'), dict) else {}
    return {
        'latency_mean': _as_num(latency.get('mean')),
        'avg_output_tps': _as_num(throughput.get('avg_output_tps')),
        'avg_req_ps': _as_num(throughput.get('avg_req_ps')),
        'input_tok_mean': _as_num(in_tok.get('mean')),
        'output_tok_mean': _as_num(out_tok.get('mean')),
        'total_tokens': _as_num(usage.get('total_tokens_count')),
        'ttft_mean': _as_num(ttft.get('mean')),
        'ttft_p90': _as_num(ttft.get('90%')),
        'ttft_p99': _as_num(ttft.get('99%')),
        'tpot_mean': _as_num(tpot.get('mean')),
        'tpot_p90': _as_num(tpot.get('90%')),
        'tpot_p99': _as_num(tpot.get('99%')),
        'n_samples': summary.get('n_samples'),
    }


def _assign_request_perf(fields: dict):
    """Unpack ``request_perf_from_summary`` into the collect_benchmark locals."""
    return (
        fields.get('latency_mean', np.nan),
        fields.get('avg_output_tps', np.nan),
        fields.get('avg_req_ps', np.nan),
        fields.get('input_tok_mean', np.nan),
        fields.get('output_tok_mean', np.nan),
        fields.get('total_tokens', np.nan),
        fields.get('ttft_mean', np.nan),
        fields.get('ttft_p90', np.nan),
        fields.get('ttft_p99', np.nan),
        fields.get('tpot_mean', np.nan),
        fields.get('tpot_p90', np.nan),
        fields.get('tpot_p99', np.nan),
    )


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))

    report_perf = request_perf_from_summary(summary0)
    backup_summary = load_backup_summary(output_dir, benchmark, model_name) if not latencies else None
    backup_perf = request_perf_from_summary(backup_summary)

    # Per-call jsonl metrics (typical MCQ / generation benches). Agent jsonl
    # usually has no perf_metrics; prefer the report's per-request summary so
    # TTFT/TPOT/latency share the same request-count口径 as ``n_samples``.
    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)
        if n_samples_unique < len(latencies):
            n_samples_unique = len(latencies)
    elif report_perf or backup_perf:
        report_n = int((summary0 or {}).get('n_samples') or 0)
        backup_n = int((backup_summary or {}).get('n_samples') or 0)
        if backup_perf is not None and backup_n > report_n:
            chosen = backup_perf
            n_samples_unique = backup_n
        else:
            chosen = report_perf or backup_perf
            if chosen.get('n_samples') is not None:
                n_samples_unique = chosen['n_samples']
        (
            latency_mean,
            avg_output_tps,
            avg_req_ps,
            input_tok_mean,
            output_tok_mean,
            total_tokens,
            ttft_mean,
            ttft_p90,
            ttft_p99,
            tpot_mean,
            tpot_p90,
            tpot_p99,
        ) = _assign_request_perf(chosen)
    else:
        # Last resort: Harbor trajectory wall-clock (no TTFT/TPOT).
        if 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 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 = np.nan
            ttft_p90 = np.nan
            ttft_p99 = np.nan
            tpot_mean = np.nan
            tpot_p90 = np.nan
            tpot_p99 = np.nan
            if n_samples_unique < len(latencies):
                n_samples_unique = len(latencies)
        else:
            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 _canonical_benchmark_label(name: Optional[str]) -> str:
    """Normalize a summary-table benchmark label (including directory aliases)."""
    try:
        if name is None or pd.isna(name):
            return ''
    except (TypeError, ValueError):
        if name is None:
            return ''
    label = str(name).strip()
    if not label or label.lower() == 'nan':
        return ''
    return BENCHMARK_NAME_ALIAS.get(label, label)


def _coerce_summary_columns(df: pd.DataFrame) -> pd.DataFrame:
    """Align an existing table to ``OUTPUT_COLUMNS``, filling missing fields."""
    out = pd.DataFrame(index=df.index, columns=OUTPUT_COLUMNS)
    for col in OUTPUT_COLUMNS:
        if col in df.columns:
            out[col] = df[col]
    return out.reset_index(drop=True)


def _summary_body(df: pd.DataFrame) -> pd.DataFrame:
    """Drop the 总计 row from a previously written summary."""
    body = _coerce_summary_columns(df)
    category = body['分类'].astype(str).str.strip()
    return body.loc[category != TOTAL_CATEGORY].reset_index(drop=True)


def _with_total_row(body: pd.DataFrame) -> pd.DataFrame:
    """Append a recomputed 总计 row. ``body`` must not already contain one."""
    if body is None or body.empty:
        return pd.DataFrame(columns=OUTPUT_COLUMNS)

    total_score = body['得分'].mean(skipna=True)
    total_time = body['实测时间(h)'].sum(skipna=True)
    total_samples = body['总样本数'].sum(skipna=True)
    total_tokens = body['累计总tokens'].sum(skipna=True)
    total_row = {col: np.nan for col in OUTPUT_COLUMNS}
    total_row.update({
        '分类': TOTAL_CATEGORY,
        'Benchmark': '',
        '得分': round(float(total_score), 4) if pd.notna(total_score) else np.nan,
        '实测时间(h)': round(float(total_time), 4) if pd.notna(total_time) else np.nan,
        '总样本数': total_samples if pd.notna(total_samples) else np.nan,
        '累计总tokens': total_tokens if pd.notna(total_tokens) else np.nan,
    })
    return pd.concat([body, pd.DataFrame([total_row], columns=OUTPUT_COLUMNS)], ignore_index=True)


def _load_existing_summary(csv_path: Path, xlsx_path: Path) -> Optional[pd.DataFrame]:
    """Load the current summary table, preferring CSV then Excel."""
    readers = (
        (csv_path, lambda p: pd.read_csv(p, encoding='utf-8-sig')),
        (xlsx_path, lambda p: pd.read_excel(p)),
    )
    for path, reader in readers:
        if not path.exists():
            continue
        try:
            df = reader(path)
        except Exception as e:
            print(f'WARNING: failed to read existing summary {path}: {e}')
            continue
        if df is None or df.empty:
            continue
        return df
    return None


def upsert_summary_rows(existing: Optional[pd.DataFrame], new_rows: list) -> pd.DataFrame:
    """Overwrite matching Benchmark rows, append unseen ones, then recompute 总计.

    Existing row order is preserved. New benchmarks are appended just before
    the 总计 row.
    """
    incoming = pd.DataFrame(new_rows, columns=OUTPUT_COLUMNS)
    incoming = incoming.drop_duplicates(subset=['Benchmark'], keep='last')
    incoming['Benchmark'] = incoming['Benchmark'].map(_canonical_benchmark_label)

    if existing is None or existing.empty:
        body_records = []
    else:
        body_records = _summary_body(existing).to_dict('records')
        for rec in body_records:
            rec['Benchmark'] = _canonical_benchmark_label(rec.get('Benchmark'))

    index_by_name = {}
    for i, rec in enumerate(body_records):
        name = str(rec.get('Benchmark') or '')
        if name:
            index_by_name[name] = i

    for row in incoming.to_dict('records'):
        name = str(row.get('Benchmark') or '')
        if not name:
            continue
        if name in index_by_name:
            body_records[index_by_name[name]] = row
        else:
            index_by_name[name] = len(body_records)
            body_records.append(row)

    body = pd.DataFrame(body_records, columns=OUTPUT_COLUMNS) if body_records else incoming
    if not body.empty:
        body = body.drop_duplicates(subset=['Benchmark'], keep='first')
    return _with_total_row(body)


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 these benchmarks are
            refreshed from disk; other rows already in the summary file are kept.
        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 upsert them into summary Excel/CSV.

    Existing rows for other benchmarks are preserved. Matching ``Benchmark``
    names are overwritten; unseen names are appended. The 总计 row is rebuilt
    from the merged table.
    """
    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

    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 = upsert_summary_rows(_load_existing_summary(csv_path, xlsx_path), rows)
    df.to_csv(csv_path, index=False, encoding='utf-8-sig')
    try:
        df.to_excel(xlsx_path, index=False)
    except Exception as e:
        print(f'WARNING: failed to write Excel {xlsx_path}: {e}')

    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()