extract_data.py

"""
extract_data.py  –  Query NATS-Bench-201 metrics for pre-weighted architectures
                    OR rank the full benchmark search space.

Usage
-----
  python extract_data.py                                        # lookup mode (default)
  python extract_data.py --mode rank --metric test-accuracy --k 10
  python extract_data.py --mode rank --metric valid-accuracy --k 5
"""

import argparse
import csv
import heapq
import sys
from pathlib import Path

from nats_bench import create

NATS_PATH   = "data/NATS-tss-v1_0-3ffb9-simple"
WEIGHTS_DIR = Path("data/weights")
HP          = 200
ALL_METRICS = ["test-accuracy", "valid-accuracy", "train-accuracy", "train-all-time"]

IMAGE_DATASETS = {
    "cifar10":        "cifar10-valid",
    "cifar100":       "cifar100",
    "ImageNet16-120": "ImageNet16-120",
}

parser = argparse.ArgumentParser(description="Query or rank NAS-Bench-201 architectures.")
parser.add_argument("--mode",   choices=["lookup", "rank"], default="lookup")
parser.add_argument("--metric", default="test-accuracy")
parser.add_argument("--k",      type=int, default=10)
args = parser.parse_args()

PBZW_INDICES = sorted(int(p.name.split(".")[0]) for p in WEIGHTS_DIR.glob("*.pickle.pbz2"))
PTH_MAP      = {int(p.name.split("-")[1]): p for p in WEIGHTS_DIR.glob("arch-*.pth")}


def hr(char="─", width=80):
    print(char * width)


def fmt(v) -> str:
    if isinstance(v, float): return f"{v:.4f}"
    if isinstance(v, int):   return str(v)
    return str(v) if v is not None else "N/A"


def print_table(rows: list[dict], col_order: list[str]):
    if not rows:
        print("  (no data)")
        return
    widths = {c: max(len(c), max(len(str(r.get(c, ""))) for r in rows)) for c in col_order}
    print("  " + "  ".join(c.ljust(widths[c]) for c in col_order))
    print("  " + "  ".join("-" * widths[c] for c in col_order))
    for r in rows:
        print("  " + "  ".join(str(r.get(c, "")).ljust(widths[c]) for c in col_order))


def fetch_metrics(idx: int, api_dataset: str, metrics: list[str]) -> dict:
    info = api.get_more_info(idx, api_dataset, hp=HP, is_random=777)
    return {m: fmt(info.get(m)) for m in metrics} if info else {m: "N/A" for m in metrics}


def sort_key(row: dict, metric: str) -> float:
    try:
        return float(row.get(metric, 0))
    except (ValueError, TypeError):
        return 0.0


def detect_pth_dataset(pth_path: Path) -> tuple[str, str]:
    """
    Each .pth file is a dict with keys: ['info', '<dataset_name>', 'all_dataset_keys'].
    Returns (dataset_name, str(all_keys)).
    """
    try:
        import torch
    except ImportError:
        return "torch-missing", "['torch module not installed']"

    try:
        obj = torch.load(pth_path, map_location="cpu", weights_only=False)
    except Exception as e:
        return "load-error", str(e)

    if not isinstance(obj, dict):
        return "unknown", f"type={type(obj).__name__}"

    keys = list(obj.keys())
    for k in keys:
        if k not in ("info", "all_dataset_keys"):
            return k, str(keys)

    return "unknown", str(keys)


def save_csv(filename: Path, fields: list[str], rows: list[dict]):
    with open(filename, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)
    print(f"  Saved {len(rows)} rows → {filename.resolve()}\n")


print(f"\nLoading NATS-Bench API  →  {NATS_PATH}")
api = create(NATS_PATH, "tss", fast_mode=True, verbose=False)
print(f"Total architectures: {len(api)}")

sample = api.get_more_info(0, "cifar10-valid", hp=HP, is_random=777)
avail  = list(sample.keys()) if sample else []
bad    = list(dict.fromkeys([m for m in ALL_METRICS + [args.metric] if m not in avail]))
if bad:
    print(f"\n[ERROR] Unknown metrics: {bad}\n  Available: {avail}")
    sys.exit(1)

print(f"Mode: {args.mode.upper()}  |  Sort metric: {args.metric}")
if args.mode == "rank":
    print(f"Top-K: {args.k}")
print()


if args.mode == "lookup":
    CSV_OUT   = Path("weights_data.csv")
    col_order = ["index"] + ALL_METRICS
    all_rows: list[dict] = []

    # 1. Image datasets (.pbz2)
    for display_name, api_name in IMAGE_DATASETS.items():
        hr()
        print(f"  {display_name.upper()}  –  {len(PBZW_INDICES)} pre-weighted architectures  (sorted by {args.metric})")
        hr()
        rows = []
        for idx in PBZW_INDICES:
            m   = fetch_metrics(idx, api_name, ALL_METRICS)
            row = {"index": idx, **m}
            rows.append(row)

        rows.sort(key=lambda r: sort_key(r, args.metric), reverse=True)
        print_table(rows, col_order)
        print()

        for row in rows:
            all_rows.append({"weight-file": f"{row['index']:06d}.pickle.pbz2", "dataset": display_name, **row})

    # 2. PyTorch datasets (.pth)
    torch_available = True
    if PTH_MAP:
        try:
            import torch
        except ImportError:
            torch_available = False
            hr()
            print("  [WARNING] Found .pth weight files but 'torch' is not installed.")
            print("            Please activate the correct conda environment (e.g. 'conda activate raml').")
            print("            Skipping DarcyFlow/NinaPro dataset detection.")
            hr()
            print()

    if PTH_MAP and torch_available:
        pth_col    = ["index", "file", "detected_by_key"] + ALL_METRICS
        by_dataset: dict[str, list[dict]] = {}

        for idx in sorted(PTH_MAP):
            pth_path          = PTH_MAP[idx]
            dataset, keys_str = detect_pth_dataset(pth_path)
            m                 = fetch_metrics(idx, "cifar10-valid", ALL_METRICS)
            row = {"index": idx, "file": pth_path.name, "detected_by_key": keys_str, **m}
            by_dataset.setdefault(dataset, []).append(row)

        for ds_name, ds_rows in sorted(by_dataset.items()):
            hr()
            print(f"  {ds_name.upper()}  –  {len(ds_rows)} .pth weight files  (sorted by {args.metric})")
            print("  Note: metrics show CIFAR-10 scores from NATS (context only).")
            hr()
            ds_rows.sort(key=lambda r: sort_key(r, args.metric), reverse=True)
            print_table(ds_rows, pth_col)
            print()

            for row in ds_rows:
                all_rows.append({
                    "weight-file": row["file"],
                    "dataset": ds_name,
                    "index": row["index"],
                    **{m: row[m] for m in ALL_METRICS}
                })

    hr("=")
    csv_fields = ["weight-file", "dataset", "index"] + ALL_METRICS
    save_csv(CSV_OUT, csv_fields, all_rows)


elif args.mode == "rank":
    CSV_OUT   = Path("ranked_data.csv")
    col_order = ["rank", "index", args.metric] + [m for m in ALL_METRICS if m != args.metric] + ["has_weights"]
    all_rows: list[dict] = []

    for display_name, api_name in IMAGE_DATASETS.items():
        hr()
        print(f"  {display_name.upper()}  –  top-{args.k} by '{args.metric}'  (hp={HP})")
        print(f"  Scanning all {len(api)} architectures …")
        hr()

        best: list[tuple[float, int]] = []
        for idx in range(len(api)):
            info  = api.get_more_info(idx, api_name, hp=HP, is_random=777)
            score = info.get(args.metric) if info else None
            if score is None:
                continue
            if len(best) < args.k:
                heapq.heappush(best, (score, idx))
            else:
                heapq.heappushpop(best, (score, idx))

        rows = []
        for rank, (score, idx) in enumerate(sorted(best, key=lambda x: -x[0]), start=1):
            info = api.get_more_info(idx, api_name, hp=HP, is_random=777)
            m    = {m: fmt(info.get(m)) for m in ALL_METRICS}
            has  = "✓ .pbz2" if idx in set(PBZW_INDICES) else "✗ missing"
            row  = {"rank": rank, "index": idx, **m, "has_weights": has}
            rows.append(row)
            all_rows.append({"dataset": display_name, **row})

        print_table(rows, col_order)
        print()

    hr("=")
    csv_fields = ["dataset", "rank", "index"] + ALL_METRICS + ["has_weights"]
    save_csv(CSV_OUT, csv_fields, all_rows)