"""
Lightweight PDF preprocessing script.

Reads PDFs from /arxiv/pdfs/{year}/, extracts structured tokens via P0's tokenizer,
and writes JSONL chunks (section/equation/figure/table/text) under exports/pdfs_structured/.

Usage:
  # Specific year range (inclusive)
  PYTHONPATH=.. python -m models.scripts.preprocess_pdfs --years 2018 2020 --max-files 1000

  # All years and all files (may be large)
  PYTHONPATH=.. python -m models.scripts.preprocess_pdfs --max-files 0

  # Use a VLM (e.g., Qwen/Qwen3-VL-2B-Instruct) for OCR
  PYTHONPATH=.. python -m models.scripts.preprocess_pdfs --max-files 0 --qwen-model Qwen/Qwen3-VL-2B-Instruct

Supports resumable runs: existing shards in the output directory are scanned and
their pdf_path entries are skipped, so re-running continues where it left off.
"""

import argparse
import json
import os
from pathlib import Path
from typing import List
import time

from models.tier3_pdf.pdf_tokenization import PDFTokenizationModel

def iter_pdf_paths(year_start: int | None, year_end: int | None, limit: int | None) -> List[Path]:
    paths: List[Path] = []
    base = Path("/arxiv/pdfs")
    years = []
    if year_start and year_end:
        years = list(range(year_start, year_end + 1))
    else:
        years = sorted([int(p.name) for p in base.iterdir() if p.is_dir() and p.name.isdigit()])
    for y in years:
        for p in (base / str(y)).glob("*.pdf"):
            paths.append(p)
            if limit and len(paths) >= limit:
                return paths
    return paths

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--years", nargs=2, type=int, help="Start/end year inclusive")
    ap.add_argument("--max-files", type=int, default=1000, help="Max PDFs to process (0 = all)")
    ap.add_argument("--out-dir", type=str, default="exports/pdfs_structured", help="Output directory for JSONL shards")
    ap.add_argument("--shard-size", type=int, default=1000, help="Records per JSONL shard")
    ap.add_argument("--progress-every", type=int, default=100, help="Print status every N PDFs")
    ap.add_argument("--resume", action="store_true", default=True, help="Skip PDFs already present in existing shards")
    ap.add_argument("--qwen-model", type=str, default=None, help="Optional Qwen VLM model name for OCR (e.g., Qwen/Qwen3-VL-2B-Instruct)")
    ap.add_argument("--max-pages", type=int, default=3, help="Max pages to parse per PDF for structured tokens")
    ap.add_argument("--no-ocr", action="store_true", help="Disable OCR (faster)")
    ap.add_argument("--no-clip", action="store_true", help="Disable CLIP embeddings (faster)")
    args = ap.parse_args()

    tok = PDFTokenizationModel()
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    year_start = args.years[0] if args.years else None
    year_end = args.years[1] if args.years else None
    limit = None if args.max_files is not None and args.max_files <= 0 else args.max_files
    paths = iter_pdf_paths(year_start, year_end, limit)

    # Resume support: gather processed pdf_paths from existing shards.
    processed = set()
    shard_idx = 0
    if args.resume:
        existing = sorted(out_dir.glob("pdf_structured_*.jsonl"))
        shard_idx = len(existing)
        for shard_file in existing:
            try:
                with shard_file.open("r", encoding="utf-8") as f:
                    for line in f:
                        try:
                            obj = json.loads(line)
                            if isinstance(obj, dict) and obj.get("pdf_path"):
                                processed.add(obj["pdf_path"])
                        except Exception:
                            continue
            except Exception:
                continue

    shard = []
    start_time = time.time()
    total = len(paths)
    for idx, pdf_path in enumerate(paths):
        if str(pdf_path) in processed:
            if (idx + 1) % args.progress_every == 0:
                elapsed = time.time() - start_time
                print(f"[resume] skipped processed {idx + 1}/{total} (elapsed {elapsed:.1f}s)")
            continue
        try:
            tokens = tok.tokenize(
                str(pdf_path),
                vlm_model=args.qwen_model,
                max_pages=args.max_pages,
                use_ocr=not args.no_ocr,
                use_clip=not args.no_clip,
            )
        except Exception as exc:
            print(f"[warn] failed to tokenize {pdf_path}: {exc}")
            continue
        shard.append({"pdf_path": str(pdf_path), "tokens": tokens})
        if len(shard) >= args.shard_size:
            out_path = out_dir / f"pdf_structured_{shard_idx:05d}.jsonl"
            with out_path.open("w", encoding="utf-8") as f:
                for rec in shard:
                    f.write(json.dumps(rec) + "\n")
            print(f"[write] shard {shard_idx} ({len(shard)} recs)")
            shard = []
            shard_idx += 1
        if (idx + 1) % args.progress_every == 0:
            elapsed = time.time() - start_time
            print(f"[status] processed {idx + 1}/{total} (elapsed {elapsed:.1f}s)")
    if shard:
        out_path = out_dir / f"pdf_structured_{shard_idx:05d}.jsonl"
        with out_path.open("w", encoding="utf-8") as f:
            for rec in shard:
                f.write(json.dumps(rec) + "\n")
        print(f"[write] shard {shard_idx} ({len(shard)} recs)")

    elapsed = time.time() - start_time
    print(f"[done] wrote shards to {out_dir} in {elapsed:.1f}s")

if __name__ == "__main__":
    main()


  # All years and all files (may be large)
  PYTHONPATH=.. python -m models.scripts.preprocess_pdfs --max-files 0

  # Use a VLM (e.g., Qwen/Qwen3-VL-2B-Instruct) for OCR
  PYTHONPATH=.. python -m models.scripts.preprocess_pdfs --max-files 0 --qwen-model Qwen/Qwen3-VL-2B-Instruct

Supports resumable runs: existing shards in the output directory are scanned and
their pdf_path entries are skipped, so re-running continues where it left off.
"""

import argparse
import json
import os
from pathlib import Path
from typing import List
import time

from models.tier3_pdf.pdf_tokenization import PDFTokenizationModel

def iter_pdf_paths(year_start: int | None, year_end: int | None, limit: int | None) -> List[Path]:
    paths: List[Path] = []
    base = Path("/arxiv/pdfs")
    years = []
    if year_start and year_end:
        years = list(range(year_start, year_end + 1))
    else:
        years = sorted([int(p.name) for p in base.iterdir() if p.is_dir() and p.name.isdigit()])
    for y in years:
        for p in (base / str(y)).glob("*.pdf"):
            paths.append(p)
            if limit and len(paths) >= limit:
                return paths
    return paths

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--years", nargs=2, type=int, help="Start/end year inclusive")
    ap.add_argument("--max-files", type=int, default=1000, help="Max PDFs to process (0 = all)")
    ap.add_argument("--out-dir", type=str, default="exports/pdfs_structured", help="Output directory for JSONL shards")
    ap.add_argument("--shard-size", type=int, default=1000, help="Records per JSONL shard")
    ap.add_argument("--progress-every", type=int, default=100, help="Print status every N PDFs")
    ap.add_argument("--resume", action="store_true", default=True, help="Skip PDFs already present in existing shards")
    ap.add_argument("--qwen-model", type=str, default=None, help="Optional Qwen VLM model name for OCR (e.g., Qwen/Qwen3-VL-2B-Instruct)")
    ap.add_argument("--max-pages", type=int, default=3, help="Max pages to parse per PDF for structured tokens")
    ap.add_argument("--no-ocr", action="store_true", help="Disable OCR (faster)")
    ap.add_argument("--no-clip", action="store_true", help="Disable CLIP embeddings (faster)")
    args = ap.parse_args()

    tok = PDFTokenizationModel()
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)


main

from models.tier3_pdf.pdf_tokenization import PDFTokenizationModel

def iter_pdf_paths(year_start: int | None, year_end: int | None, limit: int | None) -> List[Path]:
    paths: List[Path] = []
    base = Path("/arxiv/pdfs")
    years = []
    if year_start and year_end:
        years = list(range(year_start, year_end + 1))
    else:
        years = sorted([int(p.name) for p in base.iterdir() if p.is_dir() and p.name.isdigit()])
    for y in years:
        for p in (base / str(y)).glob("*.pdf"):
            paths.append(p)
            if limit and len(paths) >= limit:
                return paths
    return paths

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--years", nargs=2, type=int, help="Start/end year inclusive")
    ap.add_argument("--max-files", type=int, default=1000, help="Max PDFs to process (0 = all)")
    ap.add_argument("--out-dir", type=str, default="exports/pdfs_structured", help="Output directory for JSONL shards")
    ap.add_argument("--shard-size", type=int, default=1000, help="Records per JSONL shard")
    ap.add_argument("--progress-every", type=int, default=100, help="Print status every N PDFs")
    ap.add_argument("--resume", action="store_true", default=True, help="Skip PDFs already present in existing shards")
    ap.add_argument("--qwen-model", type=str, default=None, help="Optional Qwen VLM model name for OCR (e.g., Qwen/Qwen3-VL-2B-Instruct)")
    ap.add_argument("--max-pages", type=int, default=3, help="Max pages to parse per PDF for structured tokens")
    ap.add_argument("--no-ocr", action="store_true", help="Disable OCR (faster)")
    ap.add_argument("--no-clip", action="store_true", help="Disable CLIP embeddings (faster)")
    args = ap.parse_args()

    tok = PDFTokenizationModel()
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    year_start = args.years[0] if args.years else None
    year_end = args.years[1] if args.years else None
    limit = None if args.max_files is not None and args.max_files <= 0 else args.max_files
    paths = iter_pdf_paths(year_start, year_end, limit)

    # Resume support: gather processed pdf_paths from existing shards.
    processed = set()
    shard_idx = 0
    if args.resume:
        existing = sorted(out_dir.glob("pdf_structured_*.jsonl"))
        shard_idx = len(existing)
        for shard_file in existing:
            try:
                with shard_file.open("r", encoding="utf-8") as f:
                    for line in f:
                        try:
                            obj = json.loads(line)
                            if isinstance(obj, dict) and obj.get("pdf_path"):
                                processed.add(obj["pdf_path"])
                        except Exception:
                            continue
            except Exception:
                continue

    shard = []
    start_time = time.time()
    total = len(paths)
    for idx, pdf_path in enumerate(paths):
        if str(pdf_path) in processed:
            if (idx + 1) % args.progress_every == 0:
                elapsed = time.time() - start_time
                print(f"[resume] skipped processed {idx + 1}/{total} (elapsed {elapsed:.1f}s)")
            continue
        try:
            tokens = tok.tokenize(
                str(pdf_path),
                vlm_model=args.qwen_model,
                max_pages=args.max_pages,
                use_ocr=not args.no_ocr,
                use_clip=not args.no_clip,
            )
        except Exception as exc:
            print(f"[warn] failed to tokenize {pdf_path}: {exc}")
            continue
        shard.append({"pdf_path": str(pdf_path), "tokens": tokens})
        if len(shard) >= args.shard_size:
            out_path = out_dir / f"pdf_structured_{shard_idx:05d}.jsonl"
            with out_path.open("w", encoding="utf-8") as f:
                for rec in shard:
                    f.write(json.dumps(rec) + "\n")
            print(f"[write] shard {shard_idx} ({len(shard)} recs)")
            shard = []
            shard_idx += 1
        if (idx + 1) % args.progress_every == 0:
            elapsed = time.time() - start_time
            print(f"[status] processed {idx + 1}/{total} (elapsed {elapsed:.1f}s)")
    if shard:
        out_path = out_dir / f"pdf_structured_{shard_idx:05d}.jsonl"
        with out_path.open("w", encoding="utf-8") as f:
            for rec in shard:
                f.write(json.dumps(rec) + "\n")
        print(f"[write] shard {shard_idx} ({len(shard)} recs)")

    elapsed = time.time() - start_time
    print(f"[done] wrote shards to {out_dir} in {elapsed:.1f}s")

if __name__ == "__main__":
    main()


"""
Simple cross-domain zero-shot evaluation script for TOLBERT.

This wraps the hierarchical classification evaluator to support the scenario:

  - Train TOLBERT on domain A (e.g., CodeHierarchy).
  - Evaluate the same checkpoint on domain B (e.g., WOS) without further training.

Usage:

  # Evaluate a code-trained model on WOS spans
  python -m scripts.eval_zero_shot_cross_domain \\
      --config configs/codehierarchy_example.yaml \\
      --checkpoint checkpoints/codehierarchy/tolbert_epoch5.pt \\
      --target-config configs/wos_example.yaml \\
      --target-spans data/wos/spans_test.jsonl

The script uses:
  - base_model_name and model head structure from the *source* config
    (the one used for training),
  - but evaluation data and tokenizer settings (e.g., max_length) from
    the *target* config.
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Any, Dict

import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
import os

from tolbert.config import load_tolbert_config
from tolbert.data import TreeOfLifeDataset, collate_tree_of_life_batch
from tolbert.modeling import TOLBERT, TOLBERTConfig

def build_model(src_cfg: Dict[str, Any], checkpoint: str, device: torch.device) -> TOLBERT:
    model_cfg = TOLBERTConfig(
        base_model_name=src_cfg["base_model_name"],
        level_sizes=src_cfg["level_sizes"],
        proj_dim=src_cfg.get("proj_dim", 256),
        lambda_hier=src_cfg.get("lambda_hier", 1.0),
        lambda_path=src_cfg.get("lambda_path", 0.0),
        lambda_contrast=0.0,
    )
    model = TOLBERT(model_cfg)
    state = torch.load(checkpoint, map_location="cpu")
    model.load_state_dict(state)
    model.to(device)
    model.eval()
    return model

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Zero-shot cross-domain eval for TOLBERT.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Source (training) config for the model (defines heads).",
    )
    ap.add_argument(
        "--checkpoint",
        type=str,
        required=True,
        help="Path to trained model checkpoint on source domain.",
    )
    ap.add_argument(
        "--target-config",
        type=str,
        required=True,
        help="Target domain config (used for tokenizer and data params).",
    )
    ap.add_argument(
        "--target-spans",
        type=str,
        required=True,
        help="Spans JSONL file from the target domain (with node_path labels).",
    )
    ap.add_argument(
        "--batch-size",
        type=int,
        default=64,
        help="Evaluation batch size.",
    )
    ap.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use.",
    )
    return ap.parse_args()

def main() -> None:
    args = parse_args()

    src_cfg = load_tolbert_config(args.config)
    tgt_cfg = load_tolbert_config(args.target_config)
    device = torch.device(args.device)

    spans_path = Path(args.target_spans)
    if not spans_path.exists():
        raise FileNotFoundError(f"target_spans not found: {spans_path}")

    tokenizer = AutoTokenizer.from_pretrained(
        tgt_cfg["base_model_name"],
        cache_dir="/data/checkpoints/",  # noqa: E501
    )

    dataset = TreeOfLifeDataset(
        spans_file=str(spans_path),
        tokenizer=tokenizer,
        max_length=tgt_cfg.get("max_length", 256),
        mask_probability=tgt_cfg.get("mask_probability", 0.15),
    )
    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=tgt_cfg.get("num_workers", 0),
        collate_fn=collate_tree_of_life_batch,
    )

    model = build_model(src_cfg, checkpoint=args.checkpoint, device=device)

    level_correct: Dict[int, int] = {}
    level_total: Dict[int, int] = {}

    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            level_targets = {
                level: targets.to(device) for level, targets in batch["level_targets"].items()
            }

            out = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                level_targets=level_targets,
            )
            level_logits: Dict[str, torch.Tensor] = out["level_logits"]

            for level_int, targets in level_targets.items():
                logits = level_logits.get(str(level_int))
                if logits is None:
                    continue
                preds = logits.argmax(dim=-1)

                mask = targets != -100
                if mask.sum().item() == 0:
                    continue

                correct = (preds == targets) & mask
                num_correct = correct.sum().item()
                num_total = mask.sum().item()

                level_correct[level_int] = level_correct.get(level_int, 0) + num_correct
                level_total[level_int] = level_total.get(level_int, 0) + num_total

    print("=== Zero-shot Cross-domain Classification ===")
    for level in sorted(level_total.keys()):
        acc = level_correct[level] / max(1, level_total[level])
        print(f"Level {level}: accuracy={acc:.4f} (n={level_total[level]})")

if __name__ == "__main__":
    main()


  - but evaluation data and tokenizer settings (e.g., max_length) from
    the *target* config.
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Any, Dict

import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
import os

from tolbert.config import load_tolbert_config
from tolbert.data import TreeOfLifeDataset, collate_tree_of_life_batch
from tolbert.modeling import TOLBERT, TOLBERTConfig

def build_model(src_cfg: Dict[str, Any], checkpoint: str, device: torch.device) -> TOLBERT:
    model_cfg = TOLBERTConfig(
        base_model_name=src_cfg["base_model_name"],
        level_sizes=src_cfg["level_sizes"],
        proj_dim=src_cfg.get("proj_dim", 256),
        lambda_hier=src_cfg.get("lambda_hier", 1.0),
        lambda_path=src_cfg.get("lambda_path", 0.0),
        lambda_contrast=0.0,
    )
    model = TOLBERT(model_cfg)
    state = torch.load(checkpoint, map_location="cpu")
    model.load_state_dict(state)
    model.to(device)
    model.eval()
    return model

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Zero-shot cross-domain eval for TOLBERT.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Source (training) config for the model (defines heads).",
    )
    ap.add_argument(
        "--checkpoint",
        type=str,
        required=True,
        help="Path to trained model checkpoint on source domain.",
    )
    ap.add_argument(
        "--target-config",
        type=str,
        required=True,


from tolbert.modeling import TOLBERT, TOLBERTConfig

def build_model(src_cfg: Dict[str, Any], checkpoint: str, device: torch.device) -> TOLBERT:
    model_cfg = TOLBERTConfig(
        base_model_name=src_cfg["base_model_name"],
        level_sizes=src_cfg["level_sizes"],
        proj_dim=src_cfg.get("proj_dim", 256),
        lambda_hier=src_cfg.get("lambda_hier", 1.0),
        lambda_path=src_cfg.get("lambda_path", 0.0),
        lambda_contrast=0.0,
    )
    model = TOLBERT(model_cfg)
    state = torch.load(checkpoint, map_location="cpu")
    model.load_state_dict(state)
    model.to(device)
    model.eval()
    return model

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Zero-shot cross-domain eval for TOLBERT.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Source (training) config for the model (defines heads).",
    )
    ap.add_argument(
        "--checkpoint",
        type=str,
        required=True,
        help="Path to trained model checkpoint on source domain.",
    )
    ap.add_argument(
        "--target-config",
        type=str,
        required=True,
        help="Target domain config (used for tokenizer and data params).",
    )
    ap.add_argument(
        "--target-spans",
        type=str,
        required=True,
        help="Spans JSONL file from the target domain (with node_path labels).",
    )
    ap.add_argument(
        "--batch-size",
        type=int,
        default=64,
        help="Evaluation batch size.",
    )
    ap.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use.",
    )
    return ap.parse_args()

def main() -> None:
    args = parse_args()

    src_cfg = load_tolbert_config(args.config)
    tgt_cfg = load_tolbert_config(args.target_config)
    device = torch.device(args.device)

    spans_path = Path(args.target_spans)
    if not spans_path.exists():
        raise FileNotFoundError(f"target_spans not found: {spans_path}")

    tokenizer = AutoTokenizer.from_pretrained(
        tgt_cfg["base_model_name"],
        cache_dir="/data/checkpoints/",  # noqa: E501
    )

    dataset = TreeOfLifeDataset(
        spans_file=str(spans_path),


        "--target-spans",
        type=str,
        required=True,
        help="Spans JSONL file from the target domain (with node_path labels).",
    )
    ap.add_argument(
        "--batch-size",
        type=int,
        default=64,
        help="Evaluation batch size.",
    )
    ap.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use.",
    )
    return ap.parse_args()

def main() -> None:
    args = parse_args()

    src_cfg = load_tolbert_config(args.config)
    tgt_cfg = load_tolbert_config(args.target_config)
    device = torch.device(args.device)

    spans_path = Path(args.target_spans)
    if not spans_path.exists():
        raise FileNotFoundError(f"target_spans not found: {spans_path}")

    tokenizer = AutoTokenizer.from_pretrained(
        tgt_cfg["base_model_name"],
        cache_dir="/data/checkpoints/",  # noqa: E501
    )

    dataset = TreeOfLifeDataset(
        spans_file=str(spans_path),
        tokenizer=tokenizer,
        max_length=tgt_cfg.get("max_length", 256),
        mask_probability=tgt_cfg.get("mask_probability", 0.15),
    )
    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=tgt_cfg.get("num_workers", 0),
        collate_fn=collate_tree_of_life_batch,
    )

    model = build_model(src_cfg, checkpoint=args.checkpoint, device=device)

    level_correct: Dict[int, int] = {}
    level_total: Dict[int, int] = {}

    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            level_targets = {
                level: targets.to(device) for level, targets in batch["level_targets"].items()
            }

            out = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                level_targets=level_targets,
            )
            level_logits: Dict[str, torch.Tensor] = out["level_logits"]

            for level_int, targets in level_targets.items():
                logits = level_logits.get(str(level_int))
                if logits is None:
                    continue
                preds = logits.argmax(dim=-1)

                mask = targets != -100
                if mask.sum().item() == 0:
                    continue

                correct = (preds == targets) & mask
                num_correct = correct.sum().item()
                num_total = mask.sum().item()

                level_correct[level_int] = level_correct.get(level_int, 0) + num_correct
                level_total[level_int] = level_total.get(level_int, 0) + num_total

    print("=== Zero-shot Cross-domain Classification ===")
    for level in sorted(level_total.keys()):
        acc = level_correct[level] / max(1, level_total[level])
        print(f"Level {level}: accuracy={acc:.4f} (n={level_total[level]})")

if __name__ == "__main__":
    main()


"""
Build JSONL spans and ontology metadata for a WOS-style hierarchical text dataset.

This script is a reference implementation of the ResearchHierarchy construction
described in the TOLBERT paper. It does **not** ship any WOS data; instead, it
expects you to provide a CSV file with one row per document and three levels of
labels.

Expected input CSV schema (you can adapt this to your variant as needed):

    text_col:   column containing the text to encode (e.g., abstract)
    l1_col:     top-level field (e.g., "Computer Science")
    l2_col:     second-level subfield
    l3_col:     third-level discipline

By default we assume:

    --text-col text
    --l1-col   level1
    --l2-col   level2
    --l3-col   level3

You can override these on the command line.

Outputs:
  - --spans-out:
      JSONL file with one record per document:
        {
          "span_id": "doc_000001",
          "text": "... abstract or text ...",
          "source_id": "doc_000001",
          "node_path": [root_id, l1_id, l2_id, l3_id],
          "meta": {
             "level1": "...",
             "level2": "...",
             "level3": "..."
          }
        }

  - --nodes-out (optional):
      JSONL file of ontology nodes with fields:
        node_id, level, type, parent_id, name

  - --level-sizes-out (optional):
      Small JSON helper:
        {"level_sizes": {1: num_l1, 2: num_l2, 3: num_l3}}

These files are compatible with `tolbert.data.TreeOfLifeDataset` and the
training skeleton in `scripts/train_tolbert.py`.
"""

from __future__ import annotations

import argparse
import csv
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple

@dataclass
class DocRow:
    text: str
    l1: str
    l2: str
    l3: str

def load_wos_csv(
    path: Path,
    text_col: str,
    l1_col: str,
    l2_col: str,
    l3_col: str,
) -> List[DocRow]:
    rows: List[DocRow] = []
    with path.open("r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            text = row.get(text_col, "") or ""
            l1 = row.get(l1_col, "") or ""
            l2 = row.get(l2_col, "") or ""
            l3 = row.get(l3_col, "") or ""
            if not text:
                continue
            rows.append(DocRow(text=text, l1=l1, l2=l2, l3=l3))
    return rows

def build_ontology(rows: List[DocRow]) -> Tuple[Dict[str, int], Dict[str, int], Dict[str, int]]:
    """
    Build integer node IDs for:
      - level 1: unique l1 labels
      - level 2: unique (l1, l2) pairs
      - level 3: unique (l1, l2, l3) triples

    Node IDs are global; path uses:
      [root_id, l1_id, l2_id, l3_id]
    """
    next_id = 0
    root_id = next_id
    next_id += 1

    l1_ids: Dict[str, int] = {}
    l2_ids: Dict[str, int] = {}
    l3_ids: Dict[str, int] = {}

    for r in rows:
        if r.l1 and r.l1 not in l1_ids:
            l1_ids[r.l1] = next_id
            next_id += 1
        if r.l1 and r.l2:
            key2 = f"{r.l1}::{r.l2}"
            if key2 not in l2_ids:
                l2_ids[key2] = next_id
                next_id += 1
        if r.l1 and r.l2 and r.l3:
            key3 = f"{r.l1}::{r.l2}::{r.l3}"
            if key3 not in l3_ids:
                l3_ids[key3] = next_id
                next_id += 1

    return l1_ids, l2_ids, l3_ids

def write_nodes_jsonl(
    out_path: Path,
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
) -> None:
    with out_path.open("w", encoding="utf-8") as f:
        # Root
        f.write(
            json.dumps(
                {
                    "node_id": 0,
                    "level": 0,
                    "type": "root",
                    "parent_id": None,
                    "name": "Root",
                    "attributes": {},
                }
            )
            + "\n"
        )

        # Level 1 nodes
        for name, nid in l1_ids.items():
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 1,
                        "type": "field",
                        "parent_id": 0,
                        "name": name,
                        "attributes": {},
                    }
                )
                + "\n"
            )

        # Level 2 nodes
        for key, nid in l2_ids.items():
            l1, l2 = key.split("::", 1)
            parent_id = l1_ids[l1]
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 2,
                        "type": "subfield",
                        "parent_id": parent_id,
                        "name": l2,
                        "attributes": {"level1": l1},
                    }
                )
                + "\n"
            )

        # Level 3 nodes
        for key, nid in l3_ids.items():
            l1, l2, l3 = key.split("::", 2)
            # Parent is the (l1,l2) node
            parent_id = l2_ids[f"{l1}::{l2}"]
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 3,
                        "type": "discipline",
                        "parent_id": parent_id,
                        "name": l3,
                        "attributes": {"level1": l1, "level2": l2},
                    }
                )
                + "\n"
            )

def build_spans(
    rows: List[DocRow],
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
) -> List[Dict[str, object]]:
    spans: List[Dict[str, object]] = []
    root_id = 0
    for idx, r in enumerate(rows):
        span_id = f"doc_{idx:06d}"
        l1_id = l1_ids.get(r.l1)
        l2_id = None
        l3_id = None
        if r.l1 and r.l2:
            key2 = f"{r.l1}::{r.l2}"
            l2_id = l2_ids.get(key2)
        if r.l1 and r.l2 and r.l3:
            key3 = f"{r.l1}::{r.l2}::{r.l3}"
            l3_id = l3_ids.get(key3)

        # Build node_path; allow partially-labeled paths by omitting unknowns.
        path = [root_id]
        if l1_id is not None:
            path.append(l1_id)
        if l2_id is not None:
            path.append(l2_id)
        if l3_id is not None:
            path.append(l3_id)

        spans.append(
            {
                "span_id": span_id,
                "text": r.text,
                "source_id": span_id,
                "node_path": path,
                "meta": {
                    "level1": r.l1,
                    "level2": r.l2,
                    "level3": r.l3,
                },
            }
        )
    return spans

def write_spans_jsonl(spans: List[Dict[str, object]], out_path: Path) -> None:
    with out_path.open("w", encoding="utf-8") as f:
        for rec in spans:
            f.write(json.dumps(rec) + "\n")

def write_level_sizes(
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
    out_path: Path,
) -> None:
    level_sizes = {
        1: len(l1_ids),
        2: len(l2_ids),
        3: len(l3_ids),
    }
    out = {"level_sizes": level_sizes}
    with out_path.open("w", encoding="utf-8") as f:
        json.dump(out, f, indent=2)

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Build WOS-style spans JSONL and ontology metadata.")
    ap.add_argument("--input-csv", type=str, required=True, help="Input CSV file with text + labels.")
    ap.add_argument("--spans-out", type=str, required=True, help="Output spans JSONL file.")
    ap.add_argument("--nodes-out", type=str, default="", help="Optional nodes JSONL output.")
    ap.add_argument(
        "--level-sizes-out",
        type=str,
        default="",
        help="Optional JSON file with level_sizes helper dict.",
    )
    ap.add_argument("--text-col", type=str, default="text", help="Text column name.")
    ap.add_argument("--l1-col", type=str, default="level1", help="Level-1 label column name.")
    ap.add_argument("--l2-col", type=str, default="level2", help="Level-2 label column name.")
    ap.add_argument("--l3-col", type=str, default="level3", help="Level-3 label column name.")
    return ap.parse_args()

def main() -> None:
    args = parse_args()

    csv_path = Path(args.input_csv)
    if not csv_path.exists():
        raise FileNotFoundError(f"Input CSV not found: {csv_path}")

    rows = load_wos_csv(
        csv_path,
        text_col=args.text_col,
        l1_col=args.l1_col,
        l2_col=args.l2_col,
        l3_col=args.l3_col,
    )

    l1_ids, l2_ids, l3_ids = build_ontology(rows)

    spans = build_spans(rows, l1_ids=l1_ids, l2_ids=l2_ids, l3_ids=l3_ids)

    spans_out = Path(args.spans_out)
    spans_out.parent.mkdir(parents=True, exist_ok=True)
    write_spans_jsonl(spans, spans_out)

    if args.nodes_out:
        nodes_out = Path(args.nodes_out)
        nodes_out.parent.mkdir(parents=True, exist_ok=True)
        write_nodes_jsonl(nodes_out, l1_ids=l1_ids, l2_ids=l2_ids, l3_ids=l3_ids)

    if args.level_sizes_out:
        ls_out = Path(args.level_sizes_out)
        ls_out.parent.mkdir(parents=True, exist_ok=True)
        write_level_sizes(l1_ids=l1_ids, l2_ids=l2_ids, l3_ids=l3_ids, out_path=ls_out)

if __name__ == "__main__":
    main()


DocRow

  - --level-sizes-out (optional):
      Small JSON helper:
        {"level_sizes": {1: num_l1, 2: num_l2, 3: num_l3}}

These files are compatible with `tolbert.data.TreeOfLifeDataset` and the
training skeleton in `scripts/train_tolbert.py`.
"""

from __future__ import annotations

import argparse
import csv
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple

@dataclass
class DocRow:
    text: str
    l1: str
    l2: str
    l3: str

def load_wos_csv(
    path: Path,
    text_col: str,
    l1_col: str,
    l2_col: str,
    l3_col: str,
) -> List[DocRow]:
    rows: List[DocRow] = []
    with path.open("r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            text = row.get(text_col, "") or ""
            l1 = row.get(l1_col, "") or ""
            l2 = row.get(l2_col, "") or ""
            l3 = row.get(l3_col, "") or ""
            if not text:
                continue
            rows.append(DocRow(text=text, l1=l1, l2=l2, l3=l3))


"""

from __future__ import annotations

import argparse
import csv
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple

@dataclass
class DocRow:
    text: str
    l1: str
    l2: str
    l3: str

def load_wos_csv(
    path: Path,
    text_col: str,
    l1_col: str,
    l2_col: str,
    l3_col: str,
) -> List[DocRow]:
    rows: List[DocRow] = []
    with path.open("r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            text = row.get(text_col, "") or ""
            l1 = row.get(l1_col, "") or ""
            l2 = row.get(l2_col, "") or ""
            l3 = row.get(l3_col, "") or ""
            if not text:
                continue
            rows.append(DocRow(text=text, l1=l1, l2=l2, l3=l3))
    return rows

def build_ontology(rows: List[DocRow]) -> Tuple[Dict[str, int], Dict[str, int], Dict[str, int]]:
    """
    Build integer node IDs for:
      - level 1: unique l1 labels
      - level 2: unique (l1, l2) pairs
      - level 3: unique (l1, l2, l3) triples

    Node IDs are global; path uses:
      [root_id, l1_id, l2_id, l3_id]
    """
    next_id = 0
    root_id = next_id
    next_id += 1

    l1_ids: Dict[str, int] = {}
    l2_ids: Dict[str, int] = {}
    l3_ids: Dict[str, int] = {}


    path: Path,
    text_col: str,
    l1_col: str,
    l2_col: str,
    l3_col: str,
) -> List[DocRow]:
    rows: List[DocRow] = []
    with path.open("r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            text = row.get(text_col, "") or ""
            l1 = row.get(l1_col, "") or ""
            l2 = row.get(l2_col, "") or ""
            l3 = row.get(l3_col, "") or ""
            if not text:
                continue
            rows.append(DocRow(text=text, l1=l1, l2=l2, l3=l3))
    return rows

def build_ontology(rows: List[DocRow]) -> Tuple[Dict[str, int], Dict[str, int], Dict[str, int]]:
    """
    Build integer node IDs for:
      - level 1: unique l1 labels
      - level 2: unique (l1, l2) pairs
      - level 3: unique (l1, l2, l3) triples

    Node IDs are global; path uses:
      [root_id, l1_id, l2_id, l3_id]
    """
    next_id = 0
    root_id = next_id
    next_id += 1

    l1_ids: Dict[str, int] = {}
    l2_ids: Dict[str, int] = {}
    l3_ids: Dict[str, int] = {}

    for r in rows:
        if r.l1 and r.l1 not in l1_ids:
            l1_ids[r.l1] = next_id
            next_id += 1
        if r.l1 and r.l2:
            key2 = f"{r.l1}::{r.l2}"
            if key2 not in l2_ids:
                l2_ids[key2] = next_id
                next_id += 1
        if r.l1 and r.l2 and r.l3:
            key3 = f"{r.l1}::{r.l2}::{r.l3}"
            if key3 not in l3_ids:
                l3_ids[key3] = next_id
                next_id += 1

    return l1_ids, l2_ids, l3_ids

def write_nodes_jsonl(
    out_path: Path,
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
) -> None:
    with out_path.open("w", encoding="utf-8") as f:
        # Root
        f.write(
            json.dumps(
                {
                    "node_id": 0,
                    "level": 0,
                    "type": "root",
                    "parent_id": None,
                    "name": "Root",
                    "attributes": {},
                }


    l3_ids: Dict[str, int] = {}

    for r in rows:
        if r.l1 and r.l1 not in l1_ids:
            l1_ids[r.l1] = next_id
            next_id += 1
        if r.l1 and r.l2:
            key2 = f"{r.l1}::{r.l2}"
            if key2 not in l2_ids:
                l2_ids[key2] = next_id
                next_id += 1
        if r.l1 and r.l2 and r.l3:
            key3 = f"{r.l1}::{r.l2}::{r.l3}"
            if key3 not in l3_ids:
                l3_ids[key3] = next_id
                next_id += 1

    return l1_ids, l2_ids, l3_ids

def write_nodes_jsonl(
    out_path: Path,
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
) -> None:
    with out_path.open("w", encoding="utf-8") as f:
        # Root
        f.write(
            json.dumps(
                {
                    "node_id": 0,
                    "level": 0,
                    "type": "root",
                    "parent_id": None,
                    "name": "Root",
                    "attributes": {},
                }
            )
            + "\n"
        )

        # Level 1 nodes
        for name, nid in l1_ids.items():
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 1,
                        "type": "field",
                        "parent_id": 0,
                        "name": name,
                        "attributes": {},
                    }
                )
                + "\n"
            )

        # Level 2 nodes
        for key, nid in l2_ids.items():
            l1, l2 = key.split("::", 1)
            parent_id = l1_ids[l1]
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 2,
                        "type": "subfield",
                        "parent_id": parent_id,
                        "name": l2,
                        "attributes": {"level1": l1},
                    }
                )
                + "\n"
            )

        # Level 3 nodes
        for key, nid in l3_ids.items():
            l1, l2, l3 = key.split("::", 2)
            # Parent is the (l1,l2) node
            parent_id = l2_ids[f"{l1}::{l2}"]
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 3,
                        "type": "discipline",
                        "parent_id": parent_id,
                        "name": l3,
                        "attributes": {"level1": l1, "level2": l2},
                    }
                )
                + "\n"
            )

def build_spans(
    rows: List[DocRow],
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
) -> List[Dict[str, object]]:
    spans: List[Dict[str, object]] = []
    root_id = 0
    for idx, r in enumerate(rows):
        span_id = f"doc_{idx:06d}"
        l1_id = l1_ids.get(r.l1)
        l2_id = None
        l3_id = None
        if r.l1 and r.l2:
            key2 = f"{r.l1}::{r.l2}"
            l2_id = l2_ids.get(key2)
        if r.l1 and r.l2 and r.l3:
            key3 = f"{r.l1}::{r.l2}::{r.l3}"


        # Level 3 nodes
        for key, nid in l3_ids.items():
            l1, l2, l3 = key.split("::", 2)
            # Parent is the (l1,l2) node
            parent_id = l2_ids[f"{l1}::{l2}"]
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 3,
                        "type": "discipline",
                        "parent_id": parent_id,
                        "name": l3,
                        "attributes": {"level1": l1, "level2": l2},
                    }
                )
                + "\n"
            )

def build_spans(
    rows: List[DocRow],
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
) -> List[Dict[str, object]]:
    spans: List[Dict[str, object]] = []
    root_id = 0
    for idx, r in enumerate(rows):
        span_id = f"doc_{idx:06d}"
        l1_id = l1_ids.get(r.l1)
        l2_id = None
        l3_id = None
        if r.l1 and r.l2:
            key2 = f"{r.l1}::{r.l2}"
            l2_id = l2_ids.get(key2)
        if r.l1 and r.l2 and r.l3:
            key3 = f"{r.l1}::{r.l2}::{r.l3}"
            l3_id = l3_ids.get(key3)

        # Build node_path; allow partially-labeled paths by omitting unknowns.
        path = [root_id]
        if l1_id is not None:
            path.append(l1_id)
        if l2_id is not None:
            path.append(l2_id)
        if l3_id is not None:
            path.append(l3_id)

        spans.append(
            {
                "span_id": span_id,
                "text": r.text,
                "source_id": span_id,
                "node_path": path,
                "meta": {
                    "level1": r.l1,
                    "level2": r.l2,
                    "level3": r.l3,
                },
            }
        )
    return spans

def write_spans_jsonl(spans: List[Dict[str, object]], out_path: Path) -> None:
    with out_path.open("w", encoding="utf-8") as f:
        for rec in spans:
            f.write(json.dumps(rec) + "\n")

def write_level_sizes(
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
    out_path: Path,
) -> None:
    level_sizes = {
        1: len(l1_ids),
        2: len(l2_ids),
        3: len(l3_ids),
    }
    out = {"level_sizes": level_sizes}


            path.append(l2_id)
        if l3_id is not None:
            path.append(l3_id)

        spans.append(
            {
                "span_id": span_id,
                "text": r.text,
                "source_id": span_id,
                "node_path": path,
                "meta": {
                    "level1": r.l1,
                    "level2": r.l2,
                    "level3": r.l3,
                },
            }
        )
    return spans

def write_spans_jsonl(spans: List[Dict[str, object]], out_path: Path) -> None:
    with out_path.open("w", encoding="utf-8") as f:
        for rec in spans:
            f.write(json.dumps(rec) + "\n")

def write_level_sizes(
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
    out_path: Path,
) -> None:
    level_sizes = {
        1: len(l1_ids),
        2: len(l2_ids),
        3: len(l3_ids),
    }
    out = {"level_sizes": level_sizes}
    with out_path.open("w", encoding="utf-8") as f:
        json.dump(out, f, indent=2)

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Build WOS-style spans JSONL and ontology metadata.")


                "span_id": span_id,
                "text": r.text,
                "source_id": span_id,
                "node_path": path,
                "meta": {
                    "level1": r.l1,
                    "level2": r.l2,
                    "level3": r.l3,
                },
            }
        )
    return spans

def write_spans_jsonl(spans: List[Dict[str, object]], out_path: Path) -> None:
    with out_path.open("w", encoding="utf-8") as f:
        for rec in spans:
            f.write(json.dumps(rec) + "\n")

def write_level_sizes(
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
    out_path: Path,
) -> None:
    level_sizes = {
        1: len(l1_ids),
        2: len(l2_ids),
        3: len(l3_ids),
    }
    out = {"level_sizes": level_sizes}
    with out_path.open("w", encoding="utf-8") as f:
        json.dump(out, f, indent=2)

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Build WOS-style spans JSONL and ontology metadata.")
    ap.add_argument("--input-csv", type=str, required=True, help="Input CSV file with text + labels.")
    ap.add_argument("--spans-out", type=str, required=True, help="Output spans JSONL file.")
    ap.add_argument("--nodes-out", type=str, default="", help="Optional nodes JSONL output.")
    ap.add_argument(
        "--level-sizes-out",
        type=str,
        default="",
        help="Optional JSON file with level_sizes helper dict.",
    )
    ap.add_argument("--text-col", type=str, default="text", help="Text column name.")
    ap.add_argument("--l1-col", type=str, default="level1", help="Level-1 label column name.")
    ap.add_argument("--l2-col", type=str, default="level2", help="Level-2 label column name.")
    ap.add_argument("--l3-col", type=str, default="level3", help="Level-3 label column name.")
    return ap.parse_args()


        for rec in spans:
            f.write(json.dumps(rec) + "\n")

def write_level_sizes(
    l1_ids: Dict[str, int],
    l2_ids: Dict[str, int],
    l3_ids: Dict[str, int],
    out_path: Path,
) -> None:
    level_sizes = {
        1: len(l1_ids),
        2: len(l2_ids),
        3: len(l3_ids),
    }
    out = {"level_sizes": level_sizes}
    with out_path.open("w", encoding="utf-8") as f:
        json.dump(out, f, indent=2)

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Build WOS-style spans JSONL and ontology metadata.")
    ap.add_argument("--input-csv", type=str, required=True, help="Input CSV file with text + labels.")
    ap.add_argument("--spans-out", type=str, required=True, help="Output spans JSONL file.")
    ap.add_argument("--nodes-out", type=str, default="", help="Optional nodes JSONL output.")
    ap.add_argument(
        "--level-sizes-out",
        type=str,
        default="",
        help="Optional JSON file with level_sizes helper dict.",
    )
    ap.add_argument("--text-col", type=str, default="text", help="Text column name.")
    ap.add_argument("--l1-col", type=str, default="level1", help="Level-1 label column name.")
    ap.add_argument("--l2-col", type=str, default="level2", help="Level-2 label column name.")
    ap.add_argument("--l3-col", type=str, default="level3", help="Level-3 label column name.")
    return ap.parse_args()

def main() -> None:
    args = parse_args()

    csv_path = Path(args.input_csv)
    if not csv_path.exists():
        raise FileNotFoundError(f"Input CSV not found: {csv_path}")

    rows = load_wos_csv(
        csv_path,
        text_col=args.text_col,
        l1_col=args.l1_col,
        l2_col=args.l2_col,
        l3_col=args.l3_col,
    )

    l1_ids, l2_ids, l3_ids = build_ontology(rows)

    spans = build_spans(rows, l1_ids=l1_ids, l2_ids=l2_ids, l3_ids=l3_ids)


main

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Build WOS-style spans JSONL and ontology metadata.")
    ap.add_argument("--input-csv", type=str, required=True, help="Input CSV file with text + labels.")
    ap.add_argument("--spans-out", type=str, required=True, help="Output spans JSONL file.")
    ap.add_argument("--nodes-out", type=str, default="", help="Optional nodes JSONL output.")
    ap.add_argument(
        "--level-sizes-out",
        type=str,
        default="",
        help="Optional JSON file with level_sizes helper dict.",
    )
    ap.add_argument("--text-col", type=str, default="text", help="Text column name.")
    ap.add_argument("--l1-col", type=str, default="level1", help="Level-1 label column name.")
    ap.add_argument("--l2-col", type=str, default="level2", help="Level-2 label column name.")
    ap.add_argument("--l3-col", type=str, default="level3", help="Level-3 label column name.")
    return ap.parse_args()

def main() -> None:
    args = parse_args()

    csv_path = Path(args.input_csv)
    if not csv_path.exists():
        raise FileNotFoundError(f"Input CSV not found: {csv_path}")

    rows = load_wos_csv(
        csv_path,
        text_col=args.text_col,
        l1_col=args.l1_col,
        l2_col=args.l2_col,
        l3_col=args.l3_col,
    )

    l1_ids, l2_ids, l3_ids = build_ontology(rows)

    spans = build_spans(rows, l1_ids=l1_ids, l2_ids=l2_ids, l3_ids=l3_ids)

    spans_out = Path(args.spans_out)
    spans_out.parent.mkdir(parents=True, exist_ok=True)
    write_spans_jsonl(spans, spans_out)

    if args.nodes_out:
        nodes_out = Path(args.nodes_out)
        nodes_out.parent.mkdir(parents=True, exist_ok=True)
        write_nodes_jsonl(nodes_out, l1_ids=l1_ids, l2_ids=l2_ids, l3_ids=l3_ids)

    if args.level_sizes_out:
        ls_out = Path(args.level_sizes_out)
        ls_out.parent.mkdir(parents=True, exist_ok=True)
        write_level_sizes(l1_ids=l1_ids, l2_ids=l2_ids, l3_ids=l3_ids, out_path=ls_out)

if __name__ == "__main__":
    main()


from __future__ import annotations

import os
import ast
import hashlib
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Iterable, Set

@dataclass(frozen=True)
class FileSpan:
    file: str           # absolute path
    start_line: int     # 1-based inclusive
    end_line: int       # 1-based inclusive

@dataclass
class CGEntity:
    id: str             # stable id (e.g., fqn)
    kind: str           # module|function|class|test
    name: str
    file: str           # absolute path
    owner: Optional[str]
    start_line: int
    end_line: int

@dataclass
class CGEdge:
    src: str            # CGEntity.id
    dst: str            # CGEntity.id
    type: str           # imports|calls|owns|tests

class CodeGraph:
    def __init__(self, repo_root: str, ignore: Optional[List[str]] = None):
        self.root = os.path.abspath(repo_root)
        self.ignore_rules = [s for s in (ignore or []) if s]
        self.entities_by_id: Dict[str, CGEntity] = {}
        self.edges_list: List[CGEdge] = []
        self._file_hash: Dict[str, str] = {}
        self._id_by_module: Dict[str, str] = {}
        self._ids_by_file: Dict[str, List[str]] = {}
        self._index_identifiers: Dict[str, List[str]] = {}

    # Build
    def build(self) -> "CodeGraph":
        py_files = self._discover_py_files(self.root, self.ignore_rules)
        for abs_fp in py_files:
            mod = self._module_name_for(abs_fp)
            mid = f"py:{mod}"
            self._id_by_module[mod] = mid
            ent = CGEntity(
                id=mid, kind="module", name=mod, file=abs_fp, owner=None,
                start_line=1, end_line=self._safe_count_lines(abs_fp),
            )
            self.entities_by_id[mid] = ent
            self._ids_by_file.setdefault(abs_fp, []).append(mid)
            self._index_identifiers.setdefault(mod.lower(), []).append(mid)
            # Parse AST for defs/imports/calls
            try:
                with open(abs_fp, "r", encoding="utf-8", errors="ignore") as fh:
                    src = fh.read()
                tree = ast.parse(src)
            except Exception:
                tree = None
            if tree is None:
                continue
            # functions/classes
            for node in ast.walk(tree):
                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    name = getattr(node, "name", "")
                    a = int(getattr(node, "lineno", 1))
                    b = int(getattr(node, "end_lineno", a))
                    fid = f"py:{mod}.{name}"
                    self.entities_by_id[fid] = CGEntity(
                        id=fid, kind="function", name=name, file=abs_fp, owner=mid,
                        start_line=a, end_line=b,
                    )
                    self._ids_by_file.setdefault(abs_fp, []).append(fid)
                    self.edges_list.append(CGEdge(src=mid, dst=fid, type="owns"))
                    self._index_identifiers.setdefault(name.lower(), []).append(fid)
                elif isinstance(node, ast.ClassDef):
                    name = getattr(node, "name", "")
                    a = int(getattr(node, "lineno", 1))
                    b = int(getattr(node, "end_lineno", a))
                    cid = f"py:{mod}.{name}"
                    self.entities_by_id[cid] = CGEntity(
                        id=cid, kind="class", name=name, file=abs_fp, owner=mid,
                        start_line=a, end_line=b,
                    )
                    self._ids_by_file.setdefault(abs_fp, []).append(cid)
                    self.edges_list.append(CGEdge(src=mid, dst=cid, type="owns"))
                    self._index_identifiers.setdefault(name.lower(), []).append(cid)
            # imports (module-level)
            try:
                for node in ast.walk(tree):
                    if isinstance(node, ast.Import):
                        for alias in node.names:
                            im = alias.name
                            if not im:
                                continue
                            tgt_mod = im
                            tid = f"py:{tgt_mod}"
                            self.edges_list.append(CGEdge(src=mid, dst=tid, type="imports"))
                    elif isinstance(node, ast.ImportFrom):
                        im = node.module or ""
                        if not im:
                            continue
                        tgt_mod = im
                        tid = f"py:{tgt_mod}"
                        self.edges_list.append(CGEdge(src=mid, dst=tid, type="imports"))
            except Exception:
                pass
            # calls (best-effort): record identifiers used in Call nodes
            try:
                for node in ast.walk(tree):
                    if isinstance(node, ast.Call):
                        fn = getattr(node, "func", None)
                        name = None
                        if isinstance(fn, ast.Attribute):
                            name = getattr(fn, "attr", None)
                        elif isinstance(fn, ast.Name):
                            name = fn.id
                        if name:
                            lid = str(name).lower()
                            for cand in self._index_identifiers.get(lid, []):
                                self.edges_list.append(CGEdge(src=mid, dst=cand, type="calls"))
            except Exception:
                pass
            # tests tag
            base = os.path.basename(abs_fp)
            if base.startswith("test_") or base.endswith("_test.py"):
                self.entities_by_id[mid].kind = "test_module"
        # finalize file hashes
        self._precompute_hashes(py_files)
        return self

    # Public accessors
    def entities(self) -> Iterable[CGEntity]:
        return self.entities_by_id.values()

    def edges(self) -> Iterable[CGEdge]:
        # Filter edges whose endpoints are known (post totality)
        known = set(self.entities_by_id.keys())
        for e in self.edges_list:
            if (e.src in known) and (e.dst in known):
                yield e

    def file_hash(self, abs_path: str) -> str:
        return self._file_hash.get(abs_path) or ""

    def ids_for_file(self, abs_path: str) -> List[str]:
        return list(self._ids_by_file.get(abs_path, []))

    def find_identifier_ids(self, token: str) -> List[str]:
        return list(self._index_identifiers.get(token.lower(), []))

    # Helpers
    def _discover_py_files(self, root: str, ignore: List[str]) -> List[str]:
        out: List[str] = []
        for dirpath, dirnames, filenames in os.walk(root):
            # naive ignore: drop segments that contain any ignore pattern
            if any(ig in dirpath for ig in ignore):
                continue
            for fn in filenames:
                if not fn.endswith(".py"):
                    continue
                ap = os.path.abspath(os.path.join(dirpath, fn))
                out.append(ap)
        return out

    def _module_name_for(self, abs_file: str) -> str:
        # repo-relative without extension, path with dots
        rel = os.path.relpath(abs_file, self.root).replace("\\", "/")
        if rel.endswith(".py"):
            rel = rel[:-3]
        parts = [p for p in rel.split("/") if p and p != "__init__"]
        return ".".join(parts) or os.path.splitext(os.path.basename(abs_file))[0]

    def _safe_count_lines(self, abs_file: str) -> int:
        try:
            with open(abs_file, "r", encoding="utf-8", errors="ignore") as fh:
                return sum(1 for _ in fh)
        except Exception:
            return 1

    def _precompute_hashes(self, files: List[str]) -> None:
        for fp in files:
            try:
                with open(fp, "rb") as fh:
                    raw = fh.read()
                h = hashlib.sha256(raw).hexdigest()
            except Exception:
                h = ""
            self._file_hash[fp] = h


FileSpan

from __future__ import annotations

import os
import ast
import hashlib
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Iterable, Set

@dataclass(frozen=True)
class FileSpan:
    file: str           # absolute path
    start_line: int     # 1-based inclusive
    end_line: int       # 1-based inclusive

@dataclass
class CGEntity:
    id: str             # stable id (e.g., fqn)
    kind: str           # module|function|class|test
    name: str
    file: str           # absolute path
    owner: Optional[str]
    start_line: int
    end_line: int

@dataclass
class CGEdge:
    src: str            # CGEntity.id
    dst: str            # CGEntity.id
    type: str           # imports|calls|owns|tests


CGEntity

from __future__ import annotations

import os
import ast
import hashlib
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Iterable, Set

@dataclass(frozen=True)
class FileSpan:
    file: str           # absolute path
    start_line: int     # 1-based inclusive
    end_line: int       # 1-based inclusive

@dataclass
class CGEntity:
    id: str             # stable id (e.g., fqn)
    kind: str           # module|function|class|test
    name: str
    file: str           # absolute path
    owner: Optional[str]
    start_line: int
    end_line: int

@dataclass
class CGEdge:
    src: str            # CGEntity.id
    dst: str            # CGEntity.id
    type: str           # imports|calls|owns|tests

class CodeGraph:
    def __init__(self, repo_root: str, ignore: Optional[List[str]] = None):
        self.root = os.path.abspath(repo_root)
        self.ignore_rules = [s for s in (ignore or []) if s]
        self.entities_by_id: Dict[str, CGEntity] = {}
        self.edges_list: List[CGEdge] = []
        self._file_hash: Dict[str, str] = {}
        self._id_by_module: Dict[str, str] = {}
        self._ids_by_file: Dict[str, List[str]] = {}
        self._index_identifiers: Dict[str, List[str]] = {}


CGEdge

@dataclass(frozen=True)
class FileSpan:
    file: str           # absolute path
    start_line: int     # 1-based inclusive
    end_line: int       # 1-based inclusive

@dataclass
class CGEntity:
    id: str             # stable id (e.g., fqn)
    kind: str           # module|function|class|test
    name: str
    file: str           # absolute path
    owner: Optional[str]
    start_line: int
    end_line: int

@dataclass
class CGEdge:
    src: str            # CGEntity.id
    dst: str            # CGEntity.id
    type: str           # imports|calls|owns|tests

class CodeGraph:
    def __init__(self, repo_root: str, ignore: Optional[List[str]] = None):
        self.root = os.path.abspath(repo_root)
        self.ignore_rules = [s for s in (ignore or []) if s]
        self.entities_by_id: Dict[str, CGEntity] = {}
        self.edges_list: List[CGEdge] = []
        self._file_hash: Dict[str, str] = {}
        self._id_by_module: Dict[str, str] = {}
        self._ids_by_file: Dict[str, List[str]] = {}
        self._index_identifiers: Dict[str, List[str]] = {}

    # Build
    def build(self) -> "CodeGraph":
        py_files = self._discover_py_files(self.root, self.ignore_rules)
        for abs_fp in py_files:
            mod = self._module_name_for(abs_fp)
            mid = f"py:{mod}"
            self._id_by_module[mod] = mid


CodeGraph

@dataclass
class CGEntity:
    id: str             # stable id (e.g., fqn)
    kind: str           # module|function|class|test
    name: str
    file: str           # absolute path
    owner: Optional[str]
    start_line: int
    end_line: int

@dataclass
class CGEdge:
    src: str            # CGEntity.id
    dst: str            # CGEntity.id
    type: str           # imports|calls|owns|tests

class CodeGraph:
    def __init__(self, repo_root: str, ignore: Optional[List[str]] = None):
        self.root = os.path.abspath(repo_root)
        self.ignore_rules = [s for s in (ignore or []) if s]
        self.entities_by_id: Dict[str, CGEntity] = {}
        self.edges_list: List[CGEdge] = []
        self._file_hash: Dict[str, str] = {}
        self._id_by_module: Dict[str, str] = {}
        self._ids_by_file: Dict[str, List[str]] = {}
        self._index_identifiers: Dict[str, List[str]] = {}

    # Build
    def build(self) -> "CodeGraph":
        py_files = self._discover_py_files(self.root, self.ignore_rules)
        for abs_fp in py_files:
            mod = self._module_name_for(abs_fp)
            mid = f"py:{mod}"
            self._id_by_module[mod] = mid
            ent = CGEntity(
                id=mid, kind="module", name=mod, file=abs_fp, owner=None,
                start_line=1, end_line=self._safe_count_lines(abs_fp),
            )
            self.entities_by_id[mid] = ent
            self._ids_by_file.setdefault(abs_fp, []).append(mid)
            self._index_identifiers.setdefault(mod.lower(), []).append(mid)
            # Parse AST for defs/imports/calls
            try:
                with open(abs_fp, "r", encoding="utf-8", errors="ignore") as fh:
                    src = fh.read()
                tree = ast.parse(src)
            except Exception:
                tree = None
            if tree is None:
                continue
            # functions/classes
            for node in ast.walk(tree):
                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    name = getattr(node, "name", "")
                    a = int(getattr(node, "lineno", 1))
                    b = int(getattr(node, "end_lineno", a))
                    fid = f"py:{mod}.{name}"
                    self.entities_by_id[fid] = CGEntity(
                        id=fid, kind="function", name=name, file=abs_fp, owner=mid,
                        start_line=a, end_line=b,
                    )
                    self._ids_by_file.setdefault(abs_fp, []).append(fid)
                    self.edges_list.append(CGEdge(src=mid, dst=fid, type="owns"))
                    self._index_identifiers.setdefault(name.lower(), []).append(fid)
                elif isinstance(node, ast.ClassDef):
                    name = getattr(node, "name", "")
                    a = int(getattr(node, "lineno", 1))
                    b = int(getattr(node, "end_lineno", a))
                    cid = f"py:{mod}.{name}"
                    self.entities_by_id[cid] = CGEntity(
                        id=cid, kind="class", name=name, file=abs_fp, owner=mid,
                        start_line=a, end_line=b,
                    )
                    self._ids_by_file.setdefault(abs_fp, []).append(cid)
                    self.edges_list.append(CGEdge(src=mid, dst=cid, type="owns"))
                    self._index_identifiers.setdefault(name.lower(), []).append(cid)
            # imports (module-level)
            try:
                for node in ast.walk(tree):
                    if isinstance(node, ast.Import):
                        for alias in node.names:
                            im = alias.name
                            if not im:
                                continue
                            tgt_mod = im
                            tid = f"py:{tgt_mod}"
                            self.edges_list.append(CGEdge(src=mid, dst=tid, type="imports"))
                    elif isinstance(node, ast.ImportFrom):
                        im = node.module or ""
                        if not im:
                            continue
                        tgt_mod = im
                        tid = f"py:{tgt_mod}"
                        self.edges_list.append(CGEdge(src=mid, dst=tid, type="imports"))
            except Exception:
                pass
            # calls (best-effort): record identifiers used in Call nodes
            try:
                for node in ast.walk(tree):
                    if isinstance(node, ast.Call):
                        fn = getattr(node, "func", None)
                        name = None
                        if isinstance(fn, ast.Attribute):
                            name = getattr(fn, "attr", None)
                        elif isinstance(fn, ast.Name):
                            name = fn.id
                        if name:
                            lid = str(name).lower()
                            for cand in self._index_identifiers.get(lid, []):
                                self.edges_list.append(CGEdge(src=mid, dst=cand, type="calls"))
            except Exception:
                pass
            # tests tag
            base = os.path.basename(abs_fp)
            if base.startswith("test_") or base.endswith("_test.py"):
                self.entities_by_id[mid].kind = "test_module"
        # finalize file hashes
        self._precompute_hashes(py_files)
        return self

    # Public accessors
    def entities(self) -> Iterable[CGEntity]:
        return self.entities_by_id.values()

    def edges(self) -> Iterable[CGEdge]:
        # Filter edges whose endpoints are known (post totality)
        known = set(self.entities_by_id.keys())
        for e in self.edges_list:
            if (e.src in known) and (e.dst in known):
                yield e

    def file_hash(self, abs_path: str) -> str:
        return self._file_hash.get(abs_path) or ""

    def ids_for_file(self, abs_path: str) -> List[str]:
        return list(self._ids_by_file.get(abs_path, []))

    def find_identifier_ids(self, token: str) -> List[str]:
        return list(self._index_identifiers.get(token.lower(), []))

    # Helpers
    def _discover_py_files(self, root: str, ignore: List[str]) -> List[str]:
        out: List[str] = []
        for dirpath, dirnames, filenames in os.walk(root):
            # naive ignore: drop segments that contain any ignore pattern
            if any(ig in dirpath for ig in ignore):
                continue
            for fn in filenames:
                if not fn.endswith(".py"):
                    continue
                ap = os.path.abspath(os.path.join(dirpath, fn))
                out.append(ap)
        return out

    def _module_name_for(self, abs_file: str) -> str:
        # repo-relative without extension, path with dots
        rel = os.path.relpath(abs_file, self.root).replace("\\", "/")
        if rel.endswith(".py"):
            rel = rel[:-3]
        parts = [p for p in rel.split("/") if p and p != "__init__"]
        return ".".join(parts) or os.path.splitext(os.path.basename(abs_file))[0]

    def _safe_count_lines(self, abs_file: str) -> int:
        try:
            with open(abs_file, "r", encoding="utf-8", errors="ignore") as fh:
                return sum(1 for _ in fh)
        except Exception:
            return 1

    def _precompute_hashes(self, files: List[str]) -> None:
        for fp in files:
            try:
                with open(fp, "rb") as fh:
                    raw = fh.read()
                h = hashlib.sha256(raw).hexdigest()
            except Exception:
                h = ""
            self._file_hash[fp] = h


@dataclass
class CGEntity:
    id: str             # stable id (e.g., fqn)
    kind: str           # module|function|class|test
    name: str
    file: str           # absolute path
    owner: Optional[str]
    start_line: int
    end_line: int

@dataclass
class CGEdge:
    src: str            # CGEntity.id
    dst: str            # CGEntity.id
    type: str           # imports|calls|owns|tests

class CodeGraph:
    def __init__(self, repo_root: str, ignore: Optional[List[str]] = None):
        self.root = os.path.abspath(repo_root)
        self.ignore_rules = [s for s in (ignore or []) if s]
        self.entities_by_id: Dict[str, CGEntity] = {}
        self.edges_list: List[CGEdge] = []
        self._file_hash: Dict[str, str] = {}
        self._id_by_module: Dict[str, str] = {}
        self._ids_by_file: Dict[str, List[str]] = {}
        self._index_identifiers: Dict[str, List[str]] = {}

    # Build
    def build(self) -> "CodeGraph":
        py_files = self._discover_py_files(self.root, self.ignore_rules)
        for abs_fp in py_files:
            mod = self._module_name_for(abs_fp)
            mid = f"py:{mod}"
            self._id_by_module[mod] = mid
            ent = CGEntity(
                id=mid, kind="module", name=mod, file=abs_fp, owner=None,
                start_line=1, end_line=self._safe_count_lines(abs_fp),
            )
            self.entities_by_id[mid] = ent
            self._ids_by_file.setdefault(abs_fp, []).append(mid)
            self._index_identifiers.setdefault(mod.lower(), []).append(mid)
            # Parse AST for defs/imports/calls
            try:
                with open(abs_fp, "r", encoding="utf-8", errors="ignore") as fh:
                    src = fh.read()
                tree = ast.parse(src)


build

@dataclass
class CGEdge:
    src: str            # CGEntity.id
    dst: str            # CGEntity.id
    type: str           # imports|calls|owns|tests

class CodeGraph:
    def __init__(self, repo_root: str, ignore: Optional[List[str]] = None):
        self.root = os.path.abspath(repo_root)
        self.ignore_rules = [s for s in (ignore or []) if s]
        self.entities_by_id: Dict[str, CGEntity] = {}
        self.edges_list: List[CGEdge] = []
        self._file_hash: Dict[str, str] = {}
        self._id_by_module: Dict[str, str] = {}
        self._ids_by_file: Dict[str, List[str]] = {}
        self._index_identifiers: Dict[str, List[str]] = {}

    # Build
    def build(self) -> "CodeGraph":
        py_files = self._discover_py_files(self.root, self.ignore_rules)
        for abs_fp in py_files:
            mod = self._module_name_for(abs_fp)
            mid = f"py:{mod}"
            self._id_by_module[mod] = mid
            ent = CGEntity(
                id=mid, kind="module", name=mod, file=abs_fp, owner=None,
                start_line=1, end_line=self._safe_count_lines(abs_fp),
            )
            self.entities_by_id[mid] = ent
            self._ids_by_file.setdefault(abs_fp, []).append(mid)
            self._index_identifiers.setdefault(mod.lower(), []).append(mid)
            # Parse AST for defs/imports/calls
            try:
                with open(abs_fp, "r", encoding="utf-8", errors="ignore") as fh:
                    src = fh.read()
                tree = ast.parse(src)
            except Exception:
                tree = None
            if tree is None:
                continue
            # functions/classes
            for node in ast.walk(tree):
                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    name = getattr(node, "name", "")
                    a = int(getattr(node, "lineno", 1))
                    b = int(getattr(node, "end_lineno", a))
                    fid = f"py:{mod}.{name}"
                    self.entities_by_id[fid] = CGEntity(
                        id=fid, kind="function", name=name, file=abs_fp, owner=mid,
                        start_line=a, end_line=b,
                    )
                    self._ids_by_file.setdefault(abs_fp, []).append(fid)
                    self.edges_list.append(CGEdge(src=mid, dst=fid, type="owns"))
                    self._index_identifiers.setdefault(name.lower(), []).append(fid)
                elif isinstance(node, ast.ClassDef):
                    name = getattr(node, "name", "")
                    a = int(getattr(node, "lineno", 1))
                    b = int(getattr(node, "end_lineno", a))
                    cid = f"py:{mod}.{name}"
                    self.entities_by_id[cid] = CGEntity(
                        id=cid, kind="class", name=name, file=abs_fp, owner=mid,
                        start_line=a, end_line=b,
                    )
                    self._ids_by_file.setdefault(abs_fp, []).append(cid)
                    self.edges_list.append(CGEdge(src=mid, dst=cid, type="owns"))
                    self._index_identifiers.setdefault(name.lower(), []).append(cid)
            # imports (module-level)
            try:
                for node in ast.walk(tree):
                    if isinstance(node, ast.Import):
                        for alias in node.names:
                            im = alias.name
                            if not im:
                                continue
                            tgt_mod = im
                            tid = f"py:{tgt_mod}"
                            self.edges_list.append(CGEdge(src=mid, dst=tid, type="imports"))
                    elif isinstance(node, ast.ImportFrom):
                        im = node.module or ""
                        if not im:
                            continue
                        tgt_mod = im
                        tid = f"py:{tgt_mod}"
                        self.edges_list.append(CGEdge(src=mid, dst=tid, type="imports"))
            except Exception:
                pass
            # calls (best-effort): record identifiers used in Call nodes
            try:
                for node in ast.walk(tree):
                    if isinstance(node, ast.Call):
                        fn = getattr(node, "func", None)
                        name = None
                        if isinstance(fn, ast.Attribute):
                            name = getattr(fn, "attr", None)
                        elif isinstance(fn, ast.Name):
                            name = fn.id
                        if name:
                            lid = str(name).lower()
                            for cand in self._index_identifiers.get(lid, []):
                                self.edges_list.append(CGEdge(src=mid, dst=cand, type="calls"))
            except Exception:
                pass
            # tests tag
            base = os.path.basename(abs_fp)
            if base.startswith("test_") or base.endswith("_test.py"):
                self.entities_by_id[mid].kind = "test_module"
        # finalize file hashes
        self._precompute_hashes(py_files)
        return self

    # Public accessors
    def entities(self) -> Iterable[CGEntity]:
        return self.entities_by_id.values()

    def edges(self) -> Iterable[CGEdge]:
        # Filter edges whose endpoints are known (post totality)
        known = set(self.entities_by_id.keys())
        for e in self.edges_list:
            if (e.src in known) and (e.dst in known):
                yield e

    def file_hash(self, abs_path: str) -> str:
        return self._file_hash.get(abs_path) or ""

    def ids_for_file(self, abs_path: str) -> List[str]:
        return list(self._ids_by_file.get(abs_path, []))

    def find_identifier_ids(self, token: str) -> List[str]:
        return list(self._index_identifiers.get(token.lower(), []))


entities

                        name = None
                        if isinstance(fn, ast.Attribute):
                            name = getattr(fn, "attr", None)
                        elif isinstance(fn, ast.Name):
                            name = fn.id
                        if name:
                            lid = str(name).lower()
                            for cand in self._index_identifiers.get(lid, []):
                                self.edges_list.append(CGEdge(src=mid, dst=cand, type="calls"))
            except Exception:
                pass
            # tests tag
            base = os.path.basename(abs_fp)
            if base.startswith("test_") or base.endswith("_test.py"):
                self.entities_by_id[mid].kind = "test_module"
        # finalize file hashes
        self._precompute_hashes(py_files)
        return self

    # Public accessors
    def entities(self) -> Iterable[CGEntity]:
        return self.entities_by_id.values()

    def edges(self) -> Iterable[CGEdge]:
        # Filter edges whose endpoints are known (post totality)
        known = set(self.entities_by_id.keys())
        for e in self.edges_list:
            if (e.src in known) and (e.dst in known):
                yield e

    def file_hash(self, abs_path: str) -> str:
        return self._file_hash.get(abs_path) or ""

    def ids_for_file(self, abs_path: str) -> List[str]:
        return list(self._ids_by_file.get(abs_path, []))

    def find_identifier_ids(self, token: str) -> List[str]:
        return list(self._index_identifiers.get(token.lower(), []))

    # Helpers
    def _discover_py_files(self, root: str, ignore: List[str]) -> List[str]:
        out: List[str] = []


edges

                        elif isinstance(fn, ast.Name):
                            name = fn.id
                        if name:
                            lid = str(name).lower()
                            for cand in self._index_identifiers.get(lid, []):
                                self.edges_list.append(CGEdge(src=mid, dst=cand, type="calls"))
            except Exception:
                pass
            # tests tag
            base = os.path.basename(abs_fp)
            if base.startswith("test_") or base.endswith("_test.py"):
                self.entities_by_id[mid].kind = "test_module"
        # finalize file hashes
        self._precompute_hashes(py_files)
        return self

    # Public accessors
    def entities(self) -> Iterable[CGEntity]:
        return self.entities_by_id.values()

    def edges(self) -> Iterable[CGEdge]:
        # Filter edges whose endpoints are known (post totality)
        known = set(self.entities_by_id.keys())
        for e in self.edges_list:
            if (e.src in known) and (e.dst in known):
                yield e

    def file_hash(self, abs_path: str) -> str:
        return self._file_hash.get(abs_path) or ""

    def ids_for_file(self, abs_path: str) -> List[str]:
        return list(self._ids_by_file.get(abs_path, []))

    def find_identifier_ids(self, token: str) -> List[str]:
        return list(self._index_identifiers.get(token.lower(), []))

    # Helpers
    def _discover_py_files(self, root: str, ignore: List[str]) -> List[str]:
        out: List[str] = []
        for dirpath, dirnames, filenames in os.walk(root):
            # naive ignore: drop segments that contain any ignore pattern
            if any(ig in dirpath for ig in ignore):
                continue
            for fn in filenames:
                if not fn.endswith(".py"):
                    continue


                pass
            # tests tag
            base = os.path.basename(abs_fp)
            if base.startswith("test_") or base.endswith("_test.py"):
                self.entities_by_id[mid].kind = "test_module"
        # finalize file hashes
        self._precompute_hashes(py_files)
        return self

    # Public accessors
    def entities(self) -> Iterable[CGEntity]:
        return self.entities_by_id.values()

    def edges(self) -> Iterable[CGEdge]:
        # Filter edges whose endpoints are known (post totality)
        known = set(self.entities_by_id.keys())
        for e in self.edges_list:
            if (e.src in known) and (e.dst in known):
                yield e

    def file_hash(self, abs_path: str) -> str:
        return self._file_hash.get(abs_path) or ""

    def ids_for_file(self, abs_path: str) -> List[str]:
        return list(self._ids_by_file.get(abs_path, []))

    def find_identifier_ids(self, token: str) -> List[str]:
        return list(self._index_identifiers.get(token.lower(), []))

    # Helpers
    def _discover_py_files(self, root: str, ignore: List[str]) -> List[str]:
        out: List[str] = []
        for dirpath, dirnames, filenames in os.walk(root):
            # naive ignore: drop segments that contain any ignore pattern
            if any(ig in dirpath for ig in ignore):
                continue
            for fn in filenames:
                if not fn.endswith(".py"):
                    continue
                ap = os.path.abspath(os.path.join(dirpath, fn))
                out.append(ap)
        return out


            if base.startswith("test_") or base.endswith("_test.py"):
                self.entities_by_id[mid].kind = "test_module"
        # finalize file hashes
        self._precompute_hashes(py_files)
        return self

    # Public accessors
    def entities(self) -> Iterable[CGEntity]:
        return self.entities_by_id.values()

    def edges(self) -> Iterable[CGEdge]:
        # Filter edges whose endpoints are known (post totality)
        known = set(self.entities_by_id.keys())
        for e in self.edges_list:
            if (e.src in known) and (e.dst in known):
                yield e

    def file_hash(self, abs_path: str) -> str:
        return self._file_hash.get(abs_path) or ""

    def ids_for_file(self, abs_path: str) -> List[str]:
        return list(self._ids_by_file.get(abs_path, []))

    def find_identifier_ids(self, token: str) -> List[str]:
        return list(self._index_identifiers.get(token.lower(), []))

    # Helpers
    def _discover_py_files(self, root: str, ignore: List[str]) -> List[str]:
        out: List[str] = []
        for dirpath, dirnames, filenames in os.walk(root):
            # naive ignore: drop segments that contain any ignore pattern
            if any(ig in dirpath for ig in ignore):
                continue
            for fn in filenames:
                if not fn.endswith(".py"):
                    continue
                ap = os.path.abspath(os.path.join(dirpath, fn))
                out.append(ap)
        return out

    def _module_name_for(self, abs_file: str) -> str:
        # repo-relative without extension, path with dots


        self._precompute_hashes(py_files)
        return self

    # Public accessors
    def entities(self) -> Iterable[CGEntity]:
        return self.entities_by_id.values()

    def edges(self) -> Iterable[CGEdge]:
        # Filter edges whose endpoints are known (post totality)
        known = set(self.entities_by_id.keys())
        for e in self.edges_list:
            if (e.src in known) and (e.dst in known):
                yield e

    def file_hash(self, abs_path: str) -> str:
        return self._file_hash.get(abs_path) or ""

    def ids_for_file(self, abs_path: str) -> List[str]:
        return list(self._ids_by_file.get(abs_path, []))

    def find_identifier_ids(self, token: str) -> List[str]:
        return list(self._index_identifiers.get(token.lower(), []))

    # Helpers
    def _discover_py_files(self, root: str, ignore: List[str]) -> List[str]:
        out: List[str] = []
        for dirpath, dirnames, filenames in os.walk(root):
            # naive ignore: drop segments that contain any ignore pattern
            if any(ig in dirpath for ig in ignore):
                continue
            for fn in filenames:
                if not fn.endswith(".py"):
                    continue
                ap = os.path.abspath(os.path.join(dirpath, fn))
                out.append(ap)
        return out

    def _module_name_for(self, abs_file: str) -> str:
        # repo-relative without extension, path with dots
        rel = os.path.relpath(abs_file, self.root).replace("\\", "/")
        if rel.endswith(".py"):
            rel = rel[:-3]


    def entities(self) -> Iterable[CGEntity]:
        return self.entities_by_id.values()

    def edges(self) -> Iterable[CGEdge]:
        # Filter edges whose endpoints are known (post totality)
        known = set(self.entities_by_id.keys())
        for e in self.edges_list:
            if (e.src in known) and (e.dst in known):
                yield e

    def file_hash(self, abs_path: str) -> str:
        return self._file_hash.get(abs_path) or ""

    def ids_for_file(self, abs_path: str) -> List[str]:
        return list(self._ids_by_file.get(abs_path, []))

    def find_identifier_ids(self, token: str) -> List[str]:
        return list(self._index_identifiers.get(token.lower(), []))

    # Helpers
    def _discover_py_files(self, root: str, ignore: List[str]) -> List[str]:
        out: List[str] = []
        for dirpath, dirnames, filenames in os.walk(root):
            # naive ignore: drop segments that contain any ignore pattern
            if any(ig in dirpath for ig in ignore):
                continue
            for fn in filenames:
                if not fn.endswith(".py"):
                    continue
                ap = os.path.abspath(os.path.join(dirpath, fn))
                out.append(ap)
        return out

    def _module_name_for(self, abs_file: str) -> str:
        # repo-relative without extension, path with dots
        rel = os.path.relpath(abs_file, self.root).replace("\\", "/")
        if rel.endswith(".py"):
            rel = rel[:-3]
        parts = [p for p in rel.split("/") if p and p != "__init__"]
        return ".".join(parts) or os.path.splitext(os.path.basename(abs_file))[0]

    def _safe_count_lines(self, abs_file: str) -> int:
        try:
            with open(abs_file, "r", encoding="utf-8", errors="ignore") as fh:
                return sum(1 for _ in fh)
        except Exception:
            return 1

    def _precompute_hashes(self, files: List[str]) -> None:
        for fp in files:
            try:
                with open(fp, "rb") as fh:


    def ids_for_file(self, abs_path: str) -> List[str]:
        return list(self._ids_by_file.get(abs_path, []))

    def find_identifier_ids(self, token: str) -> List[str]:
        return list(self._index_identifiers.get(token.lower(), []))

    # Helpers
    def _discover_py_files(self, root: str, ignore: List[str]) -> List[str]:
        out: List[str] = []
        for dirpath, dirnames, filenames in os.walk(root):
            # naive ignore: drop segments that contain any ignore pattern
            if any(ig in dirpath for ig in ignore):
                continue
            for fn in filenames:
                if not fn.endswith(".py"):
                    continue
                ap = os.path.abspath(os.path.join(dirpath, fn))
                out.append(ap)
        return out

    def _module_name_for(self, abs_file: str) -> str:
        # repo-relative without extension, path with dots
        rel = os.path.relpath(abs_file, self.root).replace("\\", "/")
        if rel.endswith(".py"):
            rel = rel[:-3]
        parts = [p for p in rel.split("/") if p and p != "__init__"]
        return ".".join(parts) or os.path.splitext(os.path.basename(abs_file))[0]

    def _safe_count_lines(self, abs_file: str) -> int:
        try:
            with open(abs_file, "r", encoding="utf-8", errors="ignore") as fh:
                return sum(1 for _ in fh)
        except Exception:
            return 1

    def _precompute_hashes(self, files: List[str]) -> None:
        for fp in files:
            try:
                with open(fp, "rb") as fh:
                    raw = fh.read()
                h = hashlib.sha256(raw).hexdigest()
            except Exception:
                h = ""
            self._file_hash[fp] = h


        out: List[str] = []
        for dirpath, dirnames, filenames in os.walk(root):
            # naive ignore: drop segments that contain any ignore pattern
            if any(ig in dirpath for ig in ignore):
                continue
            for fn in filenames:
                if not fn.endswith(".py"):
                    continue
                ap = os.path.abspath(os.path.join(dirpath, fn))
                out.append(ap)
        return out

    def _module_name_for(self, abs_file: str) -> str:
        # repo-relative without extension, path with dots
        rel = os.path.relpath(abs_file, self.root).replace("\\", "/")
        if rel.endswith(".py"):
            rel = rel[:-3]
        parts = [p for p in rel.split("/") if p and p != "__init__"]
        return ".".join(parts) or os.path.splitext(os.path.basename(abs_file))[0]

    def _safe_count_lines(self, abs_file: str) -> int:
        try:
            with open(abs_file, "r", encoding="utf-8", errors="ignore") as fh:
                return sum(1 for _ in fh)
        except Exception:
            return 1

    def _precompute_hashes(self, files: List[str]) -> None:
        for fp in files:
            try:
                with open(fp, "rb") as fh:
                    raw = fh.read()
                h = hashlib.sha256(raw).hexdigest()
            except Exception:
                h = ""
            self._file_hash[fp] = h


                    continue
                ap = os.path.abspath(os.path.join(dirpath, fn))
                out.append(ap)
        return out

    def _module_name_for(self, abs_file: str) -> str:
        # repo-relative without extension, path with dots
        rel = os.path.relpath(abs_file, self.root).replace("\\", "/")
        if rel.endswith(".py"):
            rel = rel[:-3]
        parts = [p for p in rel.split("/") if p and p != "__init__"]
        return ".".join(parts) or os.path.splitext(os.path.basename(abs_file))[0]

    def _safe_count_lines(self, abs_file: str) -> int:
        try:
            with open(abs_file, "r", encoding="utf-8", errors="ignore") as fh:
                return sum(1 for _ in fh)
        except Exception:
            return 1

    def _precompute_hashes(self, files: List[str]) -> None:
        for fp in files:
            try:
                with open(fp, "rb") as fh:
                    raw = fh.read()
                h = hashlib.sha256(raw).hexdigest()
            except Exception:
                h = ""
            self._file_hash[fp] = h