scripts/pick_smoke_templates.py

"""
Pick a set of templates to use for quality-check.

Reads scripts/_template_match.csv (produced by match_templates_to_data.py),
keeps templates that match >= --min-data data files, filters by engine
(default d3-js), then picks --n templates (or all of them when --n=-1)
with a deterministic seed and writes them (and per-template list of
compatible data files) to the output plan JSON.

Two common modes:
  # smoke: 30 random templates with at least 20 data files each
  python scripts/pick_smoke_templates.py \\
      --engine d3-js --n 30 --min-data 20 --per-template-n 20 \\
      --out scripts/_smoke_templates.json

  # full d3-js: every template with at least 5 data files, up to 20 each
  python scripts/pick_smoke_templates.py \\
      --engine d3-js --n -1 --min-data 5 --per-template-n 20 \\
      --out scripts/_full_d3_plan.json
"""
import argparse
import csv
import json
import os
import random
import sys
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from config import data_resource_dirs as CFG_DATA_DIRS
from modules.chart_engine.template.template_registry import scan_templates
from modules.infographics_generator.template_utils import check_template_compatibility


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--match-csv", default="scripts/_template_match.csv")
    p.add_argument("--data", nargs="+", default=None,
                   help="One or more directories of input JSONs "
                        "(default: config.data_resource_dirs)")
    p.add_argument("--engine", default="d3-js", help="Engine to sample from")
    p.add_argument("--n", type=int, default=30,
                   help="Number of templates to sample (-1 for all qualified)")
    p.add_argument("--min-data", type=int, default=20,
                   help="Only sample templates that match at least this many data files")
    p.add_argument("--per-template-n", type=int, default=20,
                   help="Number of data files to use per template")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--out", default="scripts/_smoke_templates.json")
    return p.parse_args()


def main():
    args = parse_args()
    rng = random.Random(args.seed)

    rows = list(csv.DictReader(open(args.match_csv)))
    qualified = [r for r in rows
                 if r["engine"] == args.engine
                 and int(r["num_compatible_data"]) >= args.min_data]
    print(f"Qualified ({args.engine}, >= {args.min_data} matches): {len(qualified)}")
    if args.n < 0 or len(qualified) <= args.n:
        if args.n < 0:
            print(f"Selecting ALL {len(qualified)} qualified templates")
        else:
            print(f"WARNING: only {len(qualified)} qualified, sampling all of them")
        sample = sorted(qualified, key=lambda r: (r["chart_type"], r["chart_name"]))
    else:
        sample = rng.sample(qualified, args.n)
        sample.sort(key=lambda r: (r["chart_type"], r["chart_name"]))

    # For each chosen template, enumerate the actual data files it can use.
    print("Re-running compatibility check to enumerate per-template data files...")
    data_dirs = [Path(p) for p in (args.data or CFG_DATA_DIRS)]
    for d in data_dirs:
        if not d.is_dir():
            raise SystemExit(f"--data path is not a directory: {d}")
    json_files = []
    for d in data_dirs:
        files = sorted(d.glob("*.json"))
        json_files.extend(files)
        print(f"  {d}: {len(files)} json files")
    print(f"Total: {len(json_files)} json files across {len(data_dirs)} dir(s)")

    datas = []
    for f in json_files:
        with open(f) as fh:
            d = json.load(fh)
        d["name"] = str(f)
        rel = f"{f.parent.name}/{f.name}"
        datas.append((rel, d))

    templates = scan_templates()
    matched_files = defaultdict(list)
    for name, data in datas:
        compat = check_template_compatibility(data, templates, None)
        for tpl_key, _ in compat:
            matched_files[tpl_key].append(name)

    out_records = []
    chart_type_counter = defaultdict(int)
    for r in sample:
        tpl_key = f"{r['engine']}/{r['chart_type']}/{r['chart_name']}"
        compatible = matched_files.get(tpl_key, [])
        if len(compatible) >= args.per_template_n:
            picked = rng.sample(compatible, args.per_template_n)
        else:
            picked = list(compatible)
        out_records.append({
            "engine": r["engine"],
            "chart_type": r["chart_type"],
            "chart_name": r["chart_name"],
            "total_compatible": len(compatible),
            "picked_data_files": picked,
        })
        chart_type_counter[r["chart_type"]] += 1

    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    with open(args.out, "w") as fh:
        json.dump({
            "data_dirs": [str(d) for d in data_dirs],
            "engine": args.engine,
            "n_templates": len(out_records),
            "per_template_n": args.per_template_n,
            "templates": out_records,
        }, fh, indent=2)
    print(f"Wrote {args.out}")

    print()
    print("=== Sampled templates by chart_type ===")
    for ct, n in sorted(chart_type_counter.items(), key=lambda x: -x[1]):
        print(f"  {n:>2d}  {ct}")
    print()
    print("=== Templates and how many data files each will use ===")
    for rec in out_records:
        print(f"  {len(rec['picked_data_files']):>2d}/{rec['total_compatible']:>2d}  "
              f"{rec['chart_type']:35s}  {rec['chart_name']}")


if __name__ == "__main__":
    main()