#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import gzip
import json
from pathlib import Path
from typing import Iterable

DEFAULT_BENCH_ROOT = Path('/225040511/project/bioagent-bench')

TRUTH_FILES = {
    'alzheimer-mouse': 'pathway_comparison.csv',
    'comparative-genomics': 'cluster_annotation_mapping.csv',
    'cystic-fibrosis': 'cf_variants.csv',
    'deseq': 'up_regulated_genes.csv',
    'evolution': 'variants_shared.csv',
    'giab': 'HG001_GRCh38_1_22_v4.2.1_benchmark.vcf.gz',
    'metagenomics': 'phylum_relative_abundances.csv',
    'single-cell': 'all_clusters_de_genes.csv',
    'transcript-quant': 'truth.tsv',
    'viral-metagenomics': 'taxonomy.csv',
}

ANSWER_FILES = {
    'giab': 'answer.vcf.gz',
    'transcript-quant': 'answer.tsv',
}

FIELDNAMES = [
    'Tasks',
    'results_match',
    'Selected Tools',
    'Overhead/planning占整个流',
    'Gold Items',
    'Context Tokens',
    'Planning Latency',
    'Selection Rate',
]


def main() -> None:
    parser = argparse.ArgumentParser(description='Summarize Biomni-ReAct BioAgentBench metrics.')
    parser.add_argument('--run-root', required=True, type=Path, help='Run root produced by run_bioagent_bench_deepseek.sh')
    parser.add_argument('--bench-root', type=Path, default=DEFAULT_BENCH_ROOT)
    parser.add_argument('--output', type=Path, help='CSV output path. Default: RUN_ROOT/bioagent_bench_metrics.csv')
    args = parser.parse_args()

    rows = list(summarize(args.run_root, args.bench_root))
    output = args.output or args.run_root / 'bioagent_bench_metrics.csv'
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open('w', newline='', encoding='utf-8') as handle:
        writer = csv.DictWriter(handle, fieldnames=FIELDNAMES)
        writer.writeheader()
        writer.writerows(rows)
    print(f'Wrote {len(rows)} rows to {output}')


def summarize(run_root: Path, bench_root: Path) -> Iterable[dict[str, object]]:
    for summary_path in sorted(run_root.glob('*/run_summary.json')):
        task_id = summary_path.parent.name
        summary = json.loads(summary_path.read_text(encoding='utf-8'))
        metrics = summary.get('metrics') or {}
        gold_path = truth_path(bench_root, task_id)
        answer_path = answer_path_for(summary_path.parent, task_id)
        gold_items = count_items(gold_path) if gold_path else 0
        selected = int(metrics.get('selected_tools') or len(summary.get('selected_resources') or []))
        available = int(metrics.get('available_tools') or 0)
        selection_rate = selected / available if available else ''
        yield {
            'Tasks': task_id,
            'results_match': compare_outputs(answer_path, gold_path) if gold_path else False,
            'Selected Tools': selected,
            'Overhead/planning占整个流': round_float(metrics.get('overhead_planning_ratio')),
            'Gold Items': gold_items,
            'Context Tokens': int(metrics.get('context_tokens') or 0),
            'Planning Latency': round_float(metrics.get('planning_latency_s')),
            'Selection Rate': round_float(selection_rate),
        }


def truth_path(bench_root: Path, task_id: str) -> Path | None:
    filename = TRUTH_FILES.get(task_id)
    if filename:
        path = bench_root / 'dataset' / task_id / 'results' / filename
        if path.exists():
            return path
    results_dir = bench_root / 'dataset' / task_id / 'results'
    candidates = sorted(p for p in results_dir.glob('*') if p.is_file())
    return candidates[0] if candidates else None


def answer_path_for(workspace: Path, task_id: str) -> Path | None:
    preferred = workspace / ANSWER_FILES.get(task_id, 'answer.csv')
    if preferred.exists():
        return preferred
    candidates = []
    for pattern in ('answer.*', '*.csv', '*.tsv', '*.vcf', '*.vcf.gz'):
        candidates.extend(workspace.glob(pattern))
    candidates = [path for path in candidates if path.name not in {'retrieval_plan.json', 'run_summary.json'}]
    return sorted(set(candidates))[0] if candidates else None


def count_items(path: Path) -> int:
    if path.suffix == '.gz':
        return sum(1 for line in open_text(path) if line.strip() and not line.startswith('#'))
    if path.suffix.lower() in {'.csv', '.tsv'}:
        rows = read_table(path)
        return len(rows)
    return sum(1 for line in open_text(path) if line.strip() and not line.startswith('#'))


def compare_outputs(answer: Path | None, truth: Path | None) -> bool:
    if not answer or not truth or not answer.exists() or answer.stat().st_size == 0:
        return False
    if answer.suffix.lower() in {'.csv', '.tsv'} and truth.suffix.lower() in {'.csv', '.tsv'}:
        return normalize_table_values(answer) == normalize_table_values(truth)
    if answer.name.endswith('.vcf.gz') and truth.name.endswith('.vcf.gz'):
        return normalize_lines(answer) == normalize_lines(truth)
    return normalize_text(answer) == normalize_text(truth)


def read_table(path: Path) -> list[dict[str, str]]:
    sample = ''.join(list(open_text(path))[:5])
    delimiter = '\t' if path.suffix.lower() == '.tsv' or sample.count('\t') > sample.count(',') else ','
    with open_plain_text(path) as handle:
        reader = csv.DictReader(handle, delimiter=delimiter)
        return [{clean(key): clean(value) for key, value in row.items()} for row in reader]


def normalize_table_values(path: Path) -> list[tuple[str, ...]]:
    sample = ''.join(list(open_text(path))[:5])
    delimiter = '\t' if path.suffix.lower() == '.tsv' or sample.count('\t') > sample.count(',') else ','
    rows: list[tuple[str, ...]] = []
    with open_plain_text(path) as handle:
        for raw in handle:
            if not raw.strip():
                continue
            parts = tuple(clean(part) for part in raw.rstrip('\n').split(delimiter))
            if not rows and looks_like_header(parts):
                continue
            rows.append(parts)
    return sorted(rows)


def looks_like_header(parts: tuple[str, ...]) -> bool:
    lowered = {part.lower() for part in parts}
    known_headers = {
        'transcript_id',
        'count',
        'pathway',
        'cluster_number',
        'consensus_annotation',
        'chromosome',
        'position',
        'gene_id',
        'log2foldchange',
        'pvalue',
        'padj',
        'otu',
        'kingdom',
        'phylum',
        'cluster_id',
        'predicted_cell_type',
        'gene_name',
        'contig_count',
        'domain',
        'species',
    }
    return bool(lowered & known_headers)


def normalize_rows(rows: list[dict[str, str]]) -> list[tuple[tuple[str, str], ...]]:
    return sorted(tuple(sorted((clean(k), clean(v)) for k, v in row.items())) for row in rows)


def normalize_lines(path: Path) -> list[str]:
    return sorted(clean(line) for line in open_text(path) if line.strip() and not line.startswith('#'))


def normalize_text(path: Path) -> str:
    return '\n'.join(normalize_lines(path))


def open_text(path: Path) -> Iterable[str]:
    with open_plain_text(path) as handle:
        yield from handle


def open_plain_text(path: Path):
    if path.name.endswith('.gz'):
        return gzip.open(path, 'rt', encoding='utf-8', errors='replace')
    return path.open('r', encoding='utf-8', errors='replace', newline='')


def clean(value: object) -> str:
    if value is None:
        return ''
    return ' '.join(str(value).strip().split())


def round_float(value: object) -> object:
    if value == '':
        return ''
    try:
        return round(float(value), 6)
    except (TypeError, ValueError):
        return ''


if __name__ == '__main__':
    main()