#!/usr/bin/env python3
# coding=utf-8
"""
Export per-dataset size (n_total) from VILA registry stats JSON.

This is a lightweight post-processing step so you don't need to re-run the
2-hour fast-screen just to get sizes.

Input:
  - output/vila_stats/vila_registry_stats.json (generated by analyze_vila_registry_stats.py)

Outputs:
  - output_dir/vila_dataset_sizes.csv
  - output_dir/vila_dataset_sizes.md (top-N largest + missing)

Usage:
  python3 chengyue_script/dvlm/utils/export_vila_dataset_sizes.py \
    --stats_json output/vila_stats/vila_registry_stats.json \
    --output_dir output/vila_stats \
    --top_n 50
"""

import argparse
import csv
import json
from pathlib import Path
from typing import Any, Dict, List, Optional


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Export dataset sizes (n_total) from VILA stats JSON")
    p.add_argument(
        "--stats_json",
        type=str,
        default="output/vila_stats/vila_registry_stats.json",
        help="Path to vila_registry_stats.json",
    )
    p.add_argument(
        "--output_dir",
        type=str,
        default="output/vila_stats",
        help="Directory to write outputs",
    )
    p.add_argument(
        "--top_n",
        type=int,
        default=50,
        help="How many largest datasets to include in the markdown table",
    )
    return p.parse_args()


def main() -> int:
    args = parse_args()
    stats_path = Path(args.stats_json)
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    if not stats_path.exists():
        raise FileNotFoundError(f"stats_json not found: {stats_path}")

    with open(stats_path, "r") as f:
        root: Dict[str, Any] = json.load(f)

    datasets: Dict[str, Any] = root.get("datasets", {})
    rows: List[Dict[str, Any]] = []

    for name, info in datasets.items():
        n_total: Optional[int] = info.get("n_total")
        fatal_error: Optional[str] = info.get("fatal_error")
        n_ok = info.get("n_ok")
        n_fail = info.get("n_fail")
        fail_ratio = info.get("fail_ratio")
        rows.append(
            {
                "dataset": name,
                "n_total": n_total if isinstance(n_total, int) else "",
                "n_ok": n_ok if isinstance(n_ok, int) else "",
                "n_fail": n_fail if isinstance(n_fail, int) else "",
                "fail_ratio": f"{fail_ratio:.6f}" if isinstance(fail_ratio, (int, float)) else "",
                "fatal_error": (fatal_error or ""),
            }
        )

    # CSV (sorted by n_total desc, then name)
    def sort_key(r: Dict[str, Any]):
        n = r["n_total"]
        n_val = int(n) if isinstance(n, int) else -1
        return (-n_val, r["dataset"])

    rows_sorted = sorted(rows, key=sort_key)

    csv_path = out_dir / "vila_dataset_sizes.csv"
    with open(csv_path, "w", newline="") as f:
        w = csv.DictWriter(
            f,
            fieldnames=["dataset", "n_total", "n_ok", "n_fail", "fail_ratio", "fatal_error"],
        )
        w.writeheader()
        for r in rows_sorted:
            w.writerow(r)

    # Markdown
    ok_with_size = [r for r in rows_sorted if isinstance(r["n_total"], int) and r["n_total"] > 0]
    missing_size = [r for r in rows_sorted if not isinstance(r["n_total"], int)]

    md_path = out_dir / "vila_dataset_sizes.md"
    lines: List[str] = []
    lines.append("# VILA dataset sizes (n_total)\n")
    lines.append(f"Source: `{stats_path}`\n\n")

    lines.append(f"## Top {min(args.top_n, len(ok_with_size))} largest datasets\n")
    lines.append("|dataset|n_total|\n")
    lines.append("|---|---:|\n")
    for r in ok_with_size[: args.top_n]:
        lines.append(f"|{r['dataset']}|{r['n_total']:,}|\n")

    lines.append("\n## Datasets without n_total (likely fatal_error during load)\n")
    lines.append("|dataset|fatal_error|\n")
    lines.append("|---|---|\n")
    for r in missing_size[:200]:
        fe = r["fatal_error"]
        if len(fe) > 120:
            fe = fe[:117] + "..."
        lines.append(f"|{r['dataset']}|{fe}|\n")
    if len(missing_size) > 200:
        lines.append(f"\n... and {len(missing_size)-200} more\n")

    with open(md_path, "w") as f:
        f.writelines(lines)

    print(f"Wrote: {csv_path}")
    print(f"Wrote: {md_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())


