from tolbert.modeling import TOLBERT, TOLBERTConfig

def build_model(cfg: Dict[str, Any], checkpoint: str, device: torch.device) -> TOLBERT:
    model_cfg = TOLBERTConfig(
        base_model_name=cfg["base_model_name"],
        level_sizes=cfg["level_sizes"],
        proj_dim=cfg.get("proj_dim", 256),
        lambda_hier=cfg.get("lambda_hier", 1.0),
        lambda_path=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="Evaluate hierarchical classification for TOLBERT.")
    ap.add_argument("--config", type=str, required=True, help="Training config used for the model.")
    ap.add_argument("--checkpoint", type=str, required=True, help="Path to model .pt checkpoint.")
    ap.add_argument("--spans-file", type=str, required=True, help="Labeled spans JSONL file for eval.")
    ap.add_argument(
        "--batch-size",
        type=int,
        default=64,
        help="Evaluation batch size (default: 64).",
    )
    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()
    cfg = load_tolbert_config(args.config)
    device = torch.device(args.device)

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

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

    dataset = TreeOfLifeDataset(
        spans_file=str(spans_path),
        tokenizer=tokenizer,
        max_length=cfg.get("max_length", 256),


main

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Evaluate hierarchical classification for TOLBERT.")
    ap.add_argument("--config", type=str, required=True, help="Training config used for the model.")
    ap.add_argument("--checkpoint", type=str, required=True, help="Path to model .pt checkpoint.")
    ap.add_argument("--spans-file", type=str, required=True, help="Labeled spans JSONL file for eval.")
    ap.add_argument(
        "--batch-size",
        type=int,
        default=64,
        help="Evaluation batch size (default: 64).",
    )
    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()
    cfg = load_tolbert_config(args.config)
    device = torch.device(args.device)

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

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

    dataset = TreeOfLifeDataset(
        spans_file=str(spans_path),
        tokenizer=tokenizer,
        max_length=cfg.get("max_length", 256),
        mask_probability=cfg.get("mask_probability", 0.15),
    )

    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=cfg.get("num_workers", 0),
        collate_fn=collate_tree_of_life_batch,
    )

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

    # Metrics accumulators: per-level correct / total
    level_correct: Dict[int, int] = {}
    level_total: Dict[int, int] = {}

    # For path accuracy we track per-example correctness across supervised levels.
    path_correct = 0
    path_total = 0

    # For hierarchical precision/recall/F1 (Kiritchenko-style, single-path case),
    # we treat each (level, node_id) pair as one "label". Since every supervised
    # level has exactly one node, this reduces to counting how many levels are
    # correct across all examples.
    hier_total_true = 0
    hier_total_pred = 0
    hier_total_correct = 0

    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,
                paths=batch.get("paths"),
            )
            level_logits: Dict[str, torch.Tensor] = out["level_logits"]

            batch_size = input_ids.size(0)
            # Track per-example path correctness over all available levels.
            per_example_all_correct: List[bool] = [True] * batch_size

            for level_int, targets in level_targets.items():
                logits = level_logits[str(level_int)]
                preds = logits.argmax(dim=-1)

                # Ignore positions where target == -100 (unknown/masked level)
                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

                # For hierarchical precision/recall we count each supervised
                # (example, level) as one "true" label and one "predicted" label.
                # A label is correct iff the level prediction matches the target.
                hier_total_true += num_total
                hier_total_pred += num_total
                hier_total_correct += num_correct

                # Update per-example correctness for path accuracy
                # Only consider examples where this level is supervised.
                for i in range(batch_size):
                    if not mask[i]:
                        continue
                    if not bool(correct[i]):
                        per_example_all_correct[i] = False

            # Update path metrics: any example that had at least one supervised level
            # contributes to the denominator; it counts as correct only if all such
            # supervised levels were predicted correctly.
            for i in range(batch_size):
                # A sample is included if it has any supervised level
                has_any_level = any(
                    batch["level_targets"][lvl][i].item() != -100
                    for lvl in batch["level_targets"]
                )
                if not has_any_level:
                    continue
                path_total += 1
                if per_example_all_correct[i]:
                    path_correct += 1

    print("=== Hierarchical Classification Evaluation ===")
    for level in sorted(level_total.keys()):
        acc = level_correct[level] / max(1, level_total[level])
        # For single-label multiclass, micro-F1 equals accuracy; we report it explicitly.
        micro_f1 = acc
        print(
            f"Level {level}: "
            f"accuracy={acc:.4f} "
            f"micro_F1={micro_f1:.4f} "
            f"(n={level_total[level]})"
        )

    path_acc = path_correct / max(1, path_total)
    print(f"Path accuracy (all supervised levels correct): {path_acc:.4f} (n={path_total})")

    # Global hierarchical precision/recall/F1 over the full path, adapted from
    # Kiritchenko et al. to the single-path, single-label-per-level setting.
    # Here:
    #   - "true labels"  = all supervised (example, level) pairs,
    #   - "pred labels"  = one prediction per supervised level,
    #   - "correct"      = prediction matches the ground-truth node at that level.
    if hier_total_true > 0 and hier_total_pred > 0:
        hier_prec = hier_total_correct / hier_total_pred
        hier_rec = hier_total_correct / hier_total_true
        if hier_prec + hier_rec > 0.0:
            hier_f1 = 2.0 * hier_prec * hier_rec / (hier_prec + hier_rec)
        else:
            hier_f1 = 0.0
        print(
            "Hierarchical P/R/F1 over paths "
            f"(micro-style over all levels): "
            f"precision={hier_prec:.4f} recall={hier_rec:.4f} F1={hier_f1:.4f}"
        )
    else:
        print("Hierarchical P/R/F1 over paths: undefined (no supervised labels).")

if __name__ == "__main__":
    main()


"""
Build a simple Tree-of-Life taxonomy and span records from a single repository
using the language-agnostic RepoGraph.

This is a concrete, graph-driven implementation of the code-side part of
`docs/tree_of_life.md`:

  - Nodes:
      level 0: root
      level 1: language nodes (Python, Java, C/C++, Go, JS/TS, ...)
      level 2: repo node
      level 3: top-level directory nodes (subtrees within the repo)
      level 4: file nodes
      level 5: symbol nodes (functions / methods / classes), when available

  - Edges:
      parent_id / child_id edges forming a tree over the above nodes.

  - Spans:
      one span per file, with a node_path from root → language → repo →
      top-level directory → file.

You can use the resulting nodes.jsonl / edges.jsonl / spans.jsonl as direct
inputs to the training pipeline described in `docs/tree_of_life.md` and
`docs/training.md`.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from scripts.repo_graph import RepoGraph

def _primary_language_for_file_entity(labels: List[str]) -> Optional[str]:
    for lab in labels:
        if lab.startswith("lang:"):
            return lab.split(":", 1)[1]
    return None

def build_tree_for_repo(
    repo_root: Path,
) -> Tuple[List[Dict[str, object]], List[Dict[str, object]], List[Dict[str, object]]]:
    """
    Return (nodes, edges, spans) lists.
    """
    g = RepoGraph(str(repo_root))
    ents = list(g.entities())

    # Node id assignment
    next_id = 0

    def alloc_id() -> int:
        nonlocal next_id
        nid = next_id
        next_id += 1
        return nid

    nodes: List[Dict[str, object]] = []
    edges: List[Dict[str, object]] = []
    spans: List[Dict[str, object]] = []

    # Level 0: root
    root_id = alloc_id()
    nodes.append(
        {
            "node_id": root_id,
            "level": 0,
            "type": "root",
            "parent_id": None,
            "name": "Root",
            "attributes": {},
        }
    )

    # Language nodes (level 1)
    lang_to_id: Dict[str, int] = {}

    # Single repo node (level 2)
    repo_name = os.path.basename(os.path.abspath(repo_root))
    # We will pick a "primary language" below for the repo parent.

    # File / symbol entities from RepoGraph
    file_ents: List[Tuple[str, object]] = []  # (entity_id, entity)
    symbol_ents: List[Tuple[str, object]] = []
    for e in ents:
        if getattr(e, "kind", "") == "file":
            file_ents.append((e.id, e))
        elif getattr(e, "kind", "") in ("function", "class", "method"):
            symbol_ents.append((e.id, e))

    # Collect languages observed in files
    file_langs: Dict[str, int] = {}
    for _, e in file_ents:
        lang = _primary_language_for_file_entity(getattr(e, "labels", []))
        if not lang:
            continue
        file_langs[lang] = file_langs.get(lang, 0) + 1

    # Create language nodes
    for lang in sorted(file_langs.keys()):
        nid = alloc_id()
        lang_to_id[lang] = nid
        nodes.append(
            {
                "node_id": nid,
                "level": 1,
                "type": "language",
                "parent_id": root_id,
                "name": lang,
                "attributes": {},
            }
        )
        edges.append({"parent_id": root_id, "child_id": nid})

    # Choose a primary language for the repo (highest file count; fallback: None)
    primary_lang: Optional[str] = None
    if file_langs:
        primary_lang = max(file_langs.items(), key=lambda kv: kv[1])[0]

    repo_parent = lang_to_id.get(primary_lang, root_id)
    repo_id = alloc_id()
    nodes.append(
        {
            "node_id": repo_id,
            "level": 2,
            "type": "repo",
            "parent_id": repo_parent,
            "name": repo_name,
            "attributes": {"root": str(repo_root)},
        }
    )
    edges.append({"parent_id": repo_parent, "child_id": repo_id})

    # Level 3: top-level directories / subtrees
    subdir_to_id: Dict[str, int] = {}
    for _, e in file_ents:
        rel = str(getattr(e, "attributes", {}).get("rel_path", ""))
        top = rel.split("/", 1)[0] if "/" in rel else ""
        key = top or "(root)"
        if key in subdir_to_id:
            continue
        nid = alloc_id()
        subdir_to_id[key] = nid
        nodes.append(
            {
                "node_id": nid,
                "level": 3,
                "type": "subdir",
                "parent_id": repo_id,
                "name": key,
                "attributes": {},
            }
        )
        edges.append({"parent_id": repo_id, "child_id": nid})

    # Level 4: files
    fileid_to_nodeid: Dict[str, int] = {}
    for eid, e in file_ents:
        rel = str(getattr(e, "attributes", {}).get("rel_path", ""))
        top = rel.split("/", 1)[0] if "/" in rel else ""
        key = top or "(root)"
        parent_id = subdir_to_id[key]
        nid = alloc_id()
        fileid_to_nodeid[eid] = nid
        nodes.append(
            {
                "node_id": nid,
                "level": 4,
                "type": "file",
                "parent_id": parent_id,
                "name": rel,
                "attributes": {"artifact_uri": getattr(e, "artifact_uri", None)},
            }
        )
        edges.append({"parent_id": parent_id, "child_id": nid})

    # Level 5: symbols (functions / methods / classes).
    for _, e in symbol_ents:
        art_uri = getattr(e, "artifact_uri", None)
        if not art_uri:
            continue
        # Find the file entity that owns this symbol via artifact_uri match.
        parent_node_id: Optional[int] = None
        for fe_id, fe in file_ents:
            if getattr(fe, "artifact_uri", None) == art_uri:
                parent_node_id = fileid_to_nodeid.get(fe_id)
                break
        if parent_node_id is None:
            continue
        nid = alloc_id()
        nodes.append(
            {
                "node_id": nid,
                "level": 5,
                "type": getattr(e, "kind", "symbol"),
                "parent_id": parent_node_id,
                "name": str(getattr(e, "attributes", {}).get("name", "")),
                "attributes": {
                    "artifact_uri": art_uri,
                    "span": {
                        "start_line": getattr(getattr(e, "span", None), "start_line", None),  # type: ignore[attr-defined]  # noqa: E501
                        "end_line": getattr(getattr(e, "span", None), "end_line", None),  # type: ignore[attr-defined]  # noqa: E501
                    },
                },
            }
        )
        edges.append({"parent_id": parent_node_id, "child_id": nid})

    # Spans: one per file, rooted at file nodes.
    for eid, e in file_ents:
        rel = str(getattr(e, "attributes", {}).get("rel_path", ""))
        abs_path = repo_root / rel
        try:
            text = abs_path.read_text(encoding="utf-8")
        except Exception:
            continue
        file_node_id = fileid_to_nodeid[eid]
        # Node path: root → language (primary) → repo → subdir → file.
        lang_id = repo_parent if repo_parent != root_id else None
        top = rel.split("/", 1)[0] if "/" in rel else ""
        key = top or "(root)"
        subdir_id = subdir_to_id[key]
        path: List[int] = [root_id]
        if lang_id is not None:
            path.append(lang_id)
        path.extend([repo_id, subdir_id, file_node_id])
        spans.append(
            {
                "span_id": rel,
                "text": text,
                "source_id": rel,
                "node_path": path,
                "meta": {},
            }
        )

    return nodes, edges, spans

def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("repo_root", type=str, help="Path to a single repository root")
    ap.add_argument("--nodes-out", type=str, required=True, help="Output nodes.jsonl")
    ap.add_argument("--edges-out", type=str, required=True, help="Output edges.jsonl")
    ap.add_argument(
        "--spans-out",
        type=str,
        default=None,
        help="Optional spans.jsonl output (one span per file)",
    )
    args = ap.parse_args()

    repo_root = Path(args.repo_root)
    nodes, edges, spans = build_tree_for_repo(repo_root)

    nodes_path = Path(args.nodes_out)
    edges_path = Path(args.edges_out)
    nodes_path.parent.mkdir(parents=True, exist_ok=True)
    edges_path.parent.mkdir(parents=True, exist_ok=True)

    with nodes_path.open("w", encoding="utf-8") as f:
        for row in nodes:
            f.write(json.dumps(row) + "\n")

    with edges_path.open("w", encoding="utf-8") as f:
        for row in edges:
            f.write(json.dumps(row) + "\n")

    if args.spans_out is not None:
        spans_path = Path(args.spans_out)
        spans_path.parent.mkdir(parents=True, exist_ok=True)
        with spans_path.open("w", encoding="utf-8") as f:
            for row in spans:
                f.write(json.dumps(row) + "\n")

if __name__ == "__main__":
    main()


  - Spans:
      one span per file, with a node_path from root → language → repo →
      top-level directory → file.

You can use the resulting nodes.jsonl / edges.jsonl / spans.jsonl as direct
inputs to the training pipeline described in `docs/tree_of_life.md` and
`docs/training.md`.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from scripts.repo_graph import RepoGraph

def _primary_language_for_file_entity(labels: List[str]) -> Optional[str]:
    for lab in labels:
        if lab.startswith("lang:"):
            return lab.split(":", 1)[1]
    return None

def build_tree_for_repo(
    repo_root: Path,
) -> Tuple[List[Dict[str, object]], List[Dict[str, object]], List[Dict[str, object]]]:
    """
    Return (nodes, edges, spans) lists.
    """
    g = RepoGraph(str(repo_root))
    ents = list(g.entities())

    # Node id assignment
    next_id = 0

    def alloc_id() -> int:
        nonlocal next_id
        nid = next_id
        next_id += 1
        return nid


"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from scripts.repo_graph import RepoGraph

def _primary_language_for_file_entity(labels: List[str]) -> Optional[str]:
    for lab in labels:
        if lab.startswith("lang:"):
            return lab.split(":", 1)[1]
    return None

def build_tree_for_repo(
    repo_root: Path,
) -> Tuple[List[Dict[str, object]], List[Dict[str, object]], List[Dict[str, object]]]:
    """
    Return (nodes, edges, spans) lists.
    """
    g = RepoGraph(str(repo_root))
    ents = list(g.entities())

    # Node id assignment
    next_id = 0

    def alloc_id() -> int:
        nonlocal next_id
        nid = next_id
        next_id += 1
        return nid

    nodes: List[Dict[str, object]] = []
    edges: List[Dict[str, object]] = []
    spans: List[Dict[str, object]] = []

    # Level 0: root
    root_id = alloc_id()
    nodes.append(
        {
            "node_id": root_id,
            "level": 0,
            "type": "root",
            "parent_id": None,
            "name": "Root",
            "attributes": {},
        }
    )

    # Language nodes (level 1)
    lang_to_id: Dict[str, int] = {}

    # Single repo node (level 2)
    repo_name = os.path.basename(os.path.abspath(repo_root))
    # We will pick a "primary language" below for the repo parent.

    # File / symbol entities from RepoGraph
    file_ents: List[Tuple[str, object]] = []  # (entity_id, entity)
    symbol_ents: List[Tuple[str, object]] = []
    for e in ents:
        if getattr(e, "kind", "") == "file":
            file_ents.append((e.id, e))
        elif getattr(e, "kind", "") in ("function", "class", "method"):
            symbol_ents.append((e.id, e))

    # Collect languages observed in files
    file_langs: Dict[str, int] = {}
    for _, e in file_ents:
        lang = _primary_language_for_file_entity(getattr(e, "labels", []))
        if not lang:
            continue
        file_langs[lang] = file_langs.get(lang, 0) + 1

    # Create language nodes
    for lang in sorted(file_langs.keys()):
        nid = alloc_id()
        lang_to_id[lang] = nid
        nodes.append(
            {
                "node_id": nid,
                "level": 1,
                "type": "language",
                "parent_id": root_id,
                "name": lang,
                "attributes": {},
            }
        )
        edges.append({"parent_id": root_id, "child_id": nid})

    # Choose a primary language for the repo (highest file count; fallback: None)
    primary_lang: Optional[str] = None
    if file_langs:
        primary_lang = max(file_langs.items(), key=lambda kv: kv[1])[0]

    repo_parent = lang_to_id.get(primary_lang, root_id)
    repo_id = alloc_id()
    nodes.append(
        {
            "node_id": repo_id,
            "level": 2,
            "type": "repo",
            "parent_id": repo_parent,
            "name": repo_name,
            "attributes": {"root": str(repo_root)},
        }
    )
    edges.append({"parent_id": repo_parent, "child_id": repo_id})

    # Level 3: top-level directories / subtrees
    subdir_to_id: Dict[str, int] = {}
    for _, e in file_ents:
        rel = str(getattr(e, "attributes", {}).get("rel_path", ""))
        top = rel.split("/", 1)[0] if "/" in rel else ""
        key = top or "(root)"
        if key in subdir_to_id:
            continue
        nid = alloc_id()
        subdir_to_id[key] = nid
        nodes.append(
            {
                "node_id": nid,
                "level": 3,
                "type": "subdir",
                "parent_id": repo_id,
                "name": key,
                "attributes": {},
            }
        )
        edges.append({"parent_id": repo_id, "child_id": nid})

    # Level 4: files
    fileid_to_nodeid: Dict[str, int] = {}
    for eid, e in file_ents:
        rel = str(getattr(e, "attributes", {}).get("rel_path", ""))
        top = rel.split("/", 1)[0] if "/" in rel else ""
        key = top or "(root)"
        parent_id = subdir_to_id[key]
        nid = alloc_id()
        fileid_to_nodeid[eid] = nid
        nodes.append(
            {
                "node_id": nid,
                "level": 4,
                "type": "file",
                "parent_id": parent_id,
                "name": rel,
                "attributes": {"artifact_uri": getattr(e, "artifact_uri", None)},
            }
        )
        edges.append({"parent_id": parent_id, "child_id": nid})

    # Level 5: symbols (functions / methods / classes).
    for _, e in symbol_ents:
        art_uri = getattr(e, "artifact_uri", None)
        if not art_uri:
            continue
        # Find the file entity that owns this symbol via artifact_uri match.
        parent_node_id: Optional[int] = None
        for fe_id, fe in file_ents:
            if getattr(fe, "artifact_uri", None) == art_uri:
                parent_node_id = fileid_to_nodeid.get(fe_id)
                break
        if parent_node_id is None:
            continue
        nid = alloc_id()
        nodes.append(
            {
                "node_id": nid,
                "level": 5,
                "type": getattr(e, "kind", "symbol"),
                "parent_id": parent_node_id,
                "name": str(getattr(e, "attributes", {}).get("name", "")),
                "attributes": {
                    "artifact_uri": art_uri,
                    "span": {
                        "start_line": getattr(getattr(e, "span", None), "start_line", None),  # type: ignore[attr-defined]  # noqa: E501
                        "end_line": getattr(getattr(e, "span", None), "end_line", None),  # type: ignore[attr-defined]  # noqa: E501
                    },
                },
            }
        )
        edges.append({"parent_id": parent_node_id, "child_id": nid})

    # Spans: one per file, rooted at file nodes.
    for eid, e in file_ents:
        rel = str(getattr(e, "attributes", {}).get("rel_path", ""))
        abs_path = repo_root / rel
        try:
            text = abs_path.read_text(encoding="utf-8")
        except Exception:
            continue
        file_node_id = fileid_to_nodeid[eid]
        # Node path: root → language (primary) → repo → subdir → file.
        lang_id = repo_parent if repo_parent != root_id else None
        top = rel.split("/", 1)[0] if "/" in rel else ""
        key = top or "(root)"
        subdir_id = subdir_to_id[key]
        path: List[int] = [root_id]
        if lang_id is not None:
            path.append(lang_id)
        path.extend([repo_id, subdir_id, file_node_id])
        spans.append(
            {
                "span_id": rel,
                "text": text,
                "source_id": rel,
                "node_path": path,
                "meta": {},
            }
        )

    return nodes, edges, spans

def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("repo_root", type=str, help="Path to a single repository root")
    ap.add_argument("--nodes-out", type=str, required=True, help="Output nodes.jsonl")
    ap.add_argument("--edges-out", type=str, required=True, help="Output edges.jsonl")
    ap.add_argument(
        "--spans-out",
        type=str,
        default=None,
        help="Optional spans.jsonl output (one span per file)",
    )
    args = ap.parse_args()

    repo_root = Path(args.repo_root)
    nodes, edges, spans = build_tree_for_repo(repo_root)

    nodes_path = Path(args.nodes_out)
    edges_path = Path(args.edges_out)


        top = rel.split("/", 1)[0] if "/" in rel else ""
        key = top or "(root)"
        subdir_id = subdir_to_id[key]
        path: List[int] = [root_id]
        if lang_id is not None:
            path.append(lang_id)
        path.extend([repo_id, subdir_id, file_node_id])
        spans.append(
            {
                "span_id": rel,
                "text": text,
                "source_id": rel,
                "node_path": path,
                "meta": {},
            }
        )

    return nodes, edges, spans

def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("repo_root", type=str, help="Path to a single repository root")
    ap.add_argument("--nodes-out", type=str, required=True, help="Output nodes.jsonl")
    ap.add_argument("--edges-out", type=str, required=True, help="Output edges.jsonl")
    ap.add_argument(
        "--spans-out",
        type=str,
        default=None,
        help="Optional spans.jsonl output (one span per file)",
    )
    args = ap.parse_args()

    repo_root = Path(args.repo_root)
    nodes, edges, spans = build_tree_for_repo(repo_root)

    nodes_path = Path(args.nodes_out)
    edges_path = Path(args.edges_out)
    nodes_path.parent.mkdir(parents=True, exist_ok=True)
    edges_path.parent.mkdir(parents=True, exist_ok=True)

    with nodes_path.open("w", encoding="utf-8") as f:
        for row in nodes:
            f.write(json.dumps(row) + "\n")

    with edges_path.open("w", encoding="utf-8") as f:
        for row in edges:
            f.write(json.dumps(row) + "\n")

    if args.spans_out is not None:
        spans_path = Path(args.spans_out)
        spans_path.parent.mkdir(parents=True, exist_ok=True)
        with spans_path.open("w", encoding="utf-8") as f:
            for row in spans:
                f.write(json.dumps(row) + "\n")

if __name__ == "__main__":
    main()


def _primary_language_for_file_entity(labels: List[str]) -> Optional[str]:
    for lab in labels:
        if lab.startswith("lang:"):
            return lab.split(":", 1)[1]
    return None

def build_tree_for_repo(
    repo_root: Path,
) -> Tuple[List[Dict[str, object]], List[Dict[str, object]], List[Dict[str, object]]]:
    """
    Return (nodes, edges, spans) lists.
    """
    g = RepoGraph(str(repo_root))
    ents = list(g.entities())

    # Node id assignment
    next_id = 0

    def alloc_id() -> int:
        nonlocal next_id
        nid = next_id
        next_id += 1
        return nid

    nodes: List[Dict[str, object]] = []
    edges: List[Dict[str, object]] = []
    spans: List[Dict[str, object]] = []

    # Level 0: root
    root_id = alloc_id()
    nodes.append(
        {
            "node_id": root_id,
            "level": 0,
            "type": "root",
            "parent_id": None,
            "name": "Root",
            "attributes": {},
        }
    )

    # Language nodes (level 1)
    lang_to_id: Dict[str, int] = {}


"""
Build CodeHierarchy-style JSONL spans and simple ontology metadata.

This is a *reference* implementation of the dataset construction described
in the TOLBERT paper for the CodeHierarchy benchmark. It does not fetch or
ship any GitHub data; instead, it expects you to point it at a directory of
local repositories plus a small metadata file that assigns each repo to a
language and coarse category.

Input:
  - --repos_root:
      Directory that contains one subdirectory per repository, e.g.
        repos_root/
          repo1/
            ...
          repo2/
            ...
  - --metadata_file:
      A CSV or JSONL file with, at minimum, the fields:
        repo_name, language, category
      where:
        - repo_name matches the subdirectory name under repos_root
        - language is a string label (e.g. "Python", "Java", "C++")
        - category is a coarse repo category (e.g. "Web", "ML", "Systems")

Outputs:
  - --spans_out:
      JSONL file with one record per *file* (not per function) of the form:
        {
          "span_id": "repo/file.py",
          "text": "... full file contents ...",
          "source_id": "repo/file.py",
          "node_path": [root_id, lang_id, cat_id, repo_id],
          "meta": { ... optional extra metadata ... }
        }

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

  - --level_sizes_out (optional):
      Small JSON file with:
        {"level_sizes": [num_level0, num_level1, num_level2, num_level3]}
      which you can copy into your training config as `level_sizes`.

This script is intentionally lightweight. It is meant to make the paper's
CodeHierarchy setup reproducible given local mirrors of the repositories and
a small metadata file, without hard-coding any GitHub-specific details.
"""

from __future__ import annotations

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

@dataclass
class RepoMeta:
    repo_name: str
    language: str
    category: str

def load_metadata(path: Path) -> List[RepoMeta]:
    metas: List[RepoMeta] = []
    if path.suffix.lower() in {".jsonl", ".json"}:
        with path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                metas.append(
                    RepoMeta(
                        repo_name=obj["repo_name"],
                        language=obj["language"],
                        category=obj["category"],
                    )
                )
        return metas

    # Default: CSV with header
    with path.open("r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            metas.append(
                RepoMeta(
                    repo_name=row["repo_name"],
                    language=row["language"],
                    category=row["category"],
                )
            )
    return metas

def build_ontology(metas: List[RepoMeta]) -> Tuple[Dict[str, int], Dict[str, int], Dict[str, int]]:
    """
    Build simple integer ID mappings:
      - language -> lang_node_id
      - (language, category) -> cat_node_id
      - repo_name -> repo_node_id

    Level conventions:
      level 0: root (id 0)
      level 1: language nodes
      level 2: category nodes (per language)
      level 3: repo nodes
    """
    next_id = 0
    # Root
    root_id = next_id
    next_id += 1

    lang_ids: Dict[str, int] = {}
    cat_ids: Dict[str, int] = {}
    repo_ids: Dict[str, int] = {}

    for m in metas:
        if m.language not in lang_ids:
            lang_ids[m.language] = next_id
            next_id += 1
        cat_key = f"{m.language}::{m.category}"
        if cat_key not in cat_ids:
            cat_ids[cat_key] = next_id
            next_id += 1
        if m.repo_name not in repo_ids:
            repo_ids[m.repo_name] = next_id
            next_id += 1

    return lang_ids, cat_ids, repo_ids

def write_nodes_jsonl(
    out_path: Path,
    lang_ids: Dict[str, int],
    cat_ids: Dict[str, int],
    repo_ids: Dict[str, int],
) -> None:
    """
    Emit a minimal nodes.jsonl compatible with docs/tree_of_life.md.
    """
    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"
        )

        # Languages (level 1)
        for lang, nid in lang_ids.items():
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 1,
                        "type": "language",
                        "parent_id": 0,
                        "name": lang,
                        "attributes": {},
                    }
                )
                + "\n"
            )

        # Categories (level 2)
        for key, nid in cat_ids.items():
            lang, cat = key.split("::", 1)
            parent_id = lang_ids[lang]
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 2,
                        "type": "category",
                        "parent_id": parent_id,
                        "name": cat,
                        "attributes": {"language": lang},
                    }
                )
                + "\n"
            )

        # Repos (level 3)
        for repo, nid in repo_ids.items():
            # We do not encode parent here; it can be inferred from metadata if needed.
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 3,
                        "type": "repo",
                        "parent_id": None,
                        "name": repo,
                        "attributes": {},
                    }
                )
                + "\n"
            )

def build_span_records(
    repos_root: Path,
    metas: List[RepoMeta],
    lang_ids: Dict[str, int],
    cat_ids: Dict[str, int],
    repo_ids: Dict[str, int],
) -> List[Dict[str, object]]:
    """
    Iterate over repositories and files, generating one span per file.
    """
    root_id = 0
    spans: List[Dict[str, object]] = []

    meta_by_repo: Dict[str, RepoMeta] = {m.repo_name: m for m in metas}

    for repo_name, repo_id in repo_ids.items():
        repo_dir = repos_root / repo_name
        if not repo_dir.is_dir():
            continue

        m = meta_by_repo.get(repo_name)
        if m is None:
            continue

        lang_id = lang_ids[m.language]
        cat_key = f"{m.language}::{m.category}"
        cat_id = cat_ids[cat_key]

        for dirpath, _, filenames in os.walk(repo_dir):
            for fname in filenames:
                # Basic filter: only include source-like files
                if not fname.endswith((".py", ".java", ".cpp", ".cc", ".cxx", ".h", ".hpp")):
                    continue
                fpath = Path(dirpath) / fname
                try:
                    text = fpath.read_text(encoding="utf-8")
                except (UnicodeDecodeError, OSError):
                    continue

                rel_path = fpath.relative_to(repos_root).as_posix()
                span_id = rel_path
                node_path = [root_id, lang_id, cat_id, repo_id]

                spans.append(
                    {
                        "span_id": span_id,
                        "text": text,
                        "source_id": rel_path,
                        "node_path": node_path,
                        "meta": {
                            "repo_name": repo_name,
                            "language": m.language,
                            "category": m.category,
                        },
                    }
                )

    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(
    lang_ids: Dict[str, int],
    cat_ids: Dict[str, int],
    repo_ids: Dict[str, int],
    out_path: Path,
) -> None:
    """
    Write a small JSON helper with level_sizes for convenience.

    Note: TOLBERTConfig.level_sizes expects a dict[int, int] mapping
    level index -> number of nodes at that level (excluding the root
    level, which typically has no head). For CodeHierarchy we use:
      level 1: languages
      level 2: categories
      level 3: repos

    You can either:
      - Paste this dict into a YAML config (which preserves int keys), or
      - Load it manually and convert keys to ints before constructing
        TOLBERTConfig if you prefer JSON configs.
    """
    level_sizes = {
        1: len(lang_ids),
        2: len(cat_ids),
        3: len(repo_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 CodeHierarchy spans JSONL and ontology metadata.")
    ap.add_argument("--repos_root", type=str, required=True, help="Directory containing local repos.")
    ap.add_argument(
        "--metadata_file",
        type=str,
        required=True,
        help="CSV or JSONL file with columns: repo_name, language, category.",
    )
    ap.add_argument(
        "--spans_out",
        type=str,
        required=True,
        help="Output path for spans JSONL file.",
    )
    ap.add_argument(
        "--nodes_out",
        type=str,
        default="",
        help="Optional output path for nodes JSONL (ontology nodes).",
    )
    ap.add_argument(
        "--level_sizes_out",
        type=str,
        default="",
        help="Optional output path for level_sizes JSON file.",
    )
    return ap.parse_args()

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

    repos_root = Path(args.repos_root)
    if not repos_root.is_dir():
        raise FileNotFoundError(f"repos_root does not exist or is not a directory: {repos_root}")

    meta_path = Path(args.metadata_file)
    metas = load_metadata(meta_path)

    lang_ids, cat_ids, repo_ids = build_ontology(metas)

    spans = build_span_records(
        repos_root=repos_root,
        metas=metas,
        lang_ids=lang_ids,
        cat_ids=cat_ids,
        repo_ids=repo_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, lang_ids, cat_ids, repo_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(lang_ids, cat_ids, repo_ids, ls_out)

if __name__ == "__main__":
    main()


RepoMeta

        {"level_sizes": [num_level0, num_level1, num_level2, num_level3]}
      which you can copy into your training config as `level_sizes`.

This script is intentionally lightweight. It is meant to make the paper's
CodeHierarchy setup reproducible given local mirrors of the repositories and
a small metadata file, without hard-coding any GitHub-specific details.
"""

from __future__ import annotations

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

@dataclass
class RepoMeta:
    repo_name: str
    language: str
    category: str

def load_metadata(path: Path) -> List[RepoMeta]:
    metas: List[RepoMeta] = []
    if path.suffix.lower() in {".jsonl", ".json"}:
        with path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                metas.append(
                    RepoMeta(
                        repo_name=obj["repo_name"],
                        language=obj["language"],
                        category=obj["category"],
                    )
                )
        return metas


"""

from __future__ import annotations

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

@dataclass
class RepoMeta:
    repo_name: str
    language: str
    category: str

def load_metadata(path: Path) -> List[RepoMeta]:
    metas: List[RepoMeta] = []
    if path.suffix.lower() in {".jsonl", ".json"}:
        with path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                metas.append(
                    RepoMeta(
                        repo_name=obj["repo_name"],
                        language=obj["language"],
                        category=obj["category"],
                    )
                )
        return metas

    # Default: CSV with header
    with path.open("r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            metas.append(
                RepoMeta(
                    repo_name=row["repo_name"],
                    language=row["language"],
                    category=row["category"],
                )
            )
    return metas

def build_ontology(metas: List[RepoMeta]) -> Tuple[Dict[str, int], Dict[str, int], Dict[str, int]]:
    """
    Build simple integer ID mappings:
      - language -> lang_node_id
      - (language, category) -> cat_node_id
      - repo_name -> repo_node_id

    Level conventions:
      level 0: root (id 0)
      level 1: language nodes
      level 2: category nodes (per language)
      level 3: repo nodes
    """
    next_id = 0
    # Root
    root_id = next_id
    next_id += 1


                        language=obj["language"],
                        category=obj["category"],
                    )
                )
        return metas

    # Default: CSV with header
    with path.open("r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            metas.append(
                RepoMeta(
                    repo_name=row["repo_name"],
                    language=row["language"],
                    category=row["category"],
                )
            )
    return metas

def build_ontology(metas: List[RepoMeta]) -> Tuple[Dict[str, int], Dict[str, int], Dict[str, int]]:
    """
    Build simple integer ID mappings:
      - language -> lang_node_id
      - (language, category) -> cat_node_id
      - repo_name -> repo_node_id

    Level conventions:
      level 0: root (id 0)
      level 1: language nodes
      level 2: category nodes (per language)
      level 3: repo nodes
    """
    next_id = 0
    # Root
    root_id = next_id
    next_id += 1

    lang_ids: Dict[str, int] = {}
    cat_ids: Dict[str, int] = {}
    repo_ids: Dict[str, int] = {}

    for m in metas:
        if m.language not in lang_ids:
            lang_ids[m.language] = next_id
            next_id += 1
        cat_key = f"{m.language}::{m.category}"
        if cat_key not in cat_ids:
            cat_ids[cat_key] = next_id
            next_id += 1
        if m.repo_name not in repo_ids:
            repo_ids[m.repo_name] = next_id
            next_id += 1

    return lang_ids, cat_ids, repo_ids

def write_nodes_jsonl(
    out_path: Path,
    lang_ids: Dict[str, int],
    cat_ids: Dict[str, int],
    repo_ids: Dict[str, int],
) -> None:
    """
    Emit a minimal nodes.jsonl compatible with docs/tree_of_life.md.
    """
    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,


    lang_ids: Dict[str, int] = {}
    cat_ids: Dict[str, int] = {}
    repo_ids: Dict[str, int] = {}

    for m in metas:
        if m.language not in lang_ids:
            lang_ids[m.language] = next_id
            next_id += 1
        cat_key = f"{m.language}::{m.category}"
        if cat_key not in cat_ids:
            cat_ids[cat_key] = next_id
            next_id += 1
        if m.repo_name not in repo_ids:
            repo_ids[m.repo_name] = next_id
            next_id += 1

    return lang_ids, cat_ids, repo_ids

def write_nodes_jsonl(
    out_path: Path,
    lang_ids: Dict[str, int],
    cat_ids: Dict[str, int],
    repo_ids: Dict[str, int],
) -> None:
    """
    Emit a minimal nodes.jsonl compatible with docs/tree_of_life.md.
    """
    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"
        )

        # Languages (level 1)
        for lang, nid in lang_ids.items():
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 1,
                        "type": "language",
                        "parent_id": 0,
                        "name": lang,
                        "attributes": {},
                    }
                )
                + "\n"
            )

        # Categories (level 2)
        for key, nid in cat_ids.items():
            lang, cat = key.split("::", 1)
            parent_id = lang_ids[lang]
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 2,
                        "type": "category",
                        "parent_id": parent_id,
                        "name": cat,
                        "attributes": {"language": lang},
                    }
                )
                + "\n"
            )

        # Repos (level 3)
        for repo, nid in repo_ids.items():
            # We do not encode parent here; it can be inferred from metadata if needed.
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 3,
                        "type": "repo",
                        "parent_id": None,
                        "name": repo,
                        "attributes": {},
                    }
                )
                + "\n"
            )

def build_span_records(
    repos_root: Path,
    metas: List[RepoMeta],
    lang_ids: Dict[str, int],
    cat_ids: Dict[str, int],
    repo_ids: Dict[str, int],
) -> List[Dict[str, object]]:
    """
    Iterate over repositories and files, generating one span per file.
    """
    root_id = 0
    spans: List[Dict[str, object]] = []

    meta_by_repo: Dict[str, RepoMeta] = {m.repo_name: m for m in metas}

    for repo_name, repo_id in repo_ids.items():
        repo_dir = repos_root / repo_name
        if not repo_dir.is_dir():


            )

        # Repos (level 3)
        for repo, nid in repo_ids.items():
            # We do not encode parent here; it can be inferred from metadata if needed.
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 3,
                        "type": "repo",
                        "parent_id": None,
                        "name": repo,
                        "attributes": {},
                    }
                )
                + "\n"
            )

def build_span_records(
    repos_root: Path,
    metas: List[RepoMeta],
    lang_ids: Dict[str, int],
    cat_ids: Dict[str, int],
    repo_ids: Dict[str, int],
) -> List[Dict[str, object]]:
    """
    Iterate over repositories and files, generating one span per file.
    """
    root_id = 0
    spans: List[Dict[str, object]] = []

    meta_by_repo: Dict[str, RepoMeta] = {m.repo_name: m for m in metas}

    for repo_name, repo_id in repo_ids.items():
        repo_dir = repos_root / repo_name
        if not repo_dir.is_dir():
            continue

        m = meta_by_repo.get(repo_name)
        if m is None:
            continue

        lang_id = lang_ids[m.language]
        cat_key = f"{m.language}::{m.category}"
        cat_id = cat_ids[cat_key]

        for dirpath, _, filenames in os.walk(repo_dir):
            for fname in filenames:
                # Basic filter: only include source-like files
                if not fname.endswith((".py", ".java", ".cpp", ".cc", ".cxx", ".h", ".hpp")):
                    continue
                fpath = Path(dirpath) / fname
                try:
                    text = fpath.read_text(encoding="utf-8")
                except (UnicodeDecodeError, OSError):
                    continue

                rel_path = fpath.relative_to(repos_root).as_posix()
                span_id = rel_path
                node_path = [root_id, lang_id, cat_id, repo_id]

                spans.append(
                    {
                        "span_id": span_id,
                        "text": text,
                        "source_id": rel_path,
                        "node_path": node_path,
                        "meta": {
                            "repo_name": repo_name,
                            "language": m.language,
                            "category": m.category,
                        },
                    }
                )

    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(
    lang_ids: Dict[str, int],
    cat_ids: Dict[str, int],
    repo_ids: Dict[str, int],
    out_path: Path,
) -> None:
    """
    Write a small JSON helper with level_sizes for convenience.

    Note: TOLBERTConfig.level_sizes expects a dict[int, int] mapping
    level index -> number of nodes at that level (excluding the root
    level, which typically has no head). For CodeHierarchy we use:


                span_id = rel_path
                node_path = [root_id, lang_id, cat_id, repo_id]

                spans.append(
                    {
                        "span_id": span_id,
                        "text": text,
                        "source_id": rel_path,
                        "node_path": node_path,
                        "meta": {
                            "repo_name": repo_name,
                            "language": m.language,
                            "category": m.category,
                        },
                    }
                )

    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(
    lang_ids: Dict[str, int],
    cat_ids: Dict[str, int],
    repo_ids: Dict[str, int],
    out_path: Path,
) -> None:
    """
    Write a small JSON helper with level_sizes for convenience.

    Note: TOLBERTConfig.level_sizes expects a dict[int, int] mapping
    level index -> number of nodes at that level (excluding the root
    level, which typically has no head). For CodeHierarchy we use:
      level 1: languages
      level 2: categories
      level 3: repos

    You can either:
      - Paste this dict into a YAML config (which preserves int keys), or


                        "text": text,
                        "source_id": rel_path,
                        "node_path": node_path,
                        "meta": {
                            "repo_name": repo_name,
                            "language": m.language,
                            "category": m.category,
                        },
                    }
                )

    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(
    lang_ids: Dict[str, int],
    cat_ids: Dict[str, int],
    repo_ids: Dict[str, int],
    out_path: Path,
) -> None:
    """
    Write a small JSON helper with level_sizes for convenience.

    Note: TOLBERTConfig.level_sizes expects a dict[int, int] mapping
    level index -> number of nodes at that level (excluding the root
    level, which typically has no head). For CodeHierarchy we use:
      level 1: languages
      level 2: categories
      level 3: repos

    You can either:
      - Paste this dict into a YAML config (which preserves int keys), or
      - Load it manually and convert keys to ints before constructing
        TOLBERTConfig if you prefer JSON configs.
    """
    level_sizes = {
        1: len(lang_ids),
        2: len(cat_ids),
        3: len(repo_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 CodeHierarchy spans JSONL and ontology metadata.")
    ap.add_argument("--repos_root", type=str, required=True, help="Directory containing local repos.")
    ap.add_argument(
        "--metadata_file",
        type=str,
        required=True,
        help="CSV or JSONL file with columns: repo_name, language, category.",
    )
    ap.add_argument(
        "--spans_out",
        type=str,
        required=True,
        help="Output path for spans JSONL file.",
    )
    ap.add_argument(
        "--nodes_out",
        type=str,


    level, which typically has no head). For CodeHierarchy we use:
      level 1: languages
      level 2: categories
      level 3: repos

    You can either:
      - Paste this dict into a YAML config (which preserves int keys), or
      - Load it manually and convert keys to ints before constructing
        TOLBERTConfig if you prefer JSON configs.
    """
    level_sizes = {
        1: len(lang_ids),
        2: len(cat_ids),
        3: len(repo_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 CodeHierarchy spans JSONL and ontology metadata.")
    ap.add_argument("--repos_root", type=str, required=True, help="Directory containing local repos.")
    ap.add_argument(
        "--metadata_file",
        type=str,
        required=True,
        help="CSV or JSONL file with columns: repo_name, language, category.",
    )
    ap.add_argument(
        "--spans_out",
        type=str,
        required=True,
        help="Output path for spans JSONL file.",
    )
    ap.add_argument(
        "--nodes_out",
        type=str,
        default="",
        help="Optional output path for nodes JSONL (ontology nodes).",
    )
    ap.add_argument(
        "--level_sizes_out",
        type=str,
        default="",
        help="Optional output path for level_sizes JSON file.",
    )
    return ap.parse_args()

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

    repos_root = Path(args.repos_root)
    if not repos_root.is_dir():
        raise FileNotFoundError(f"repos_root does not exist or is not a directory: {repos_root}")

    meta_path = Path(args.metadata_file)
    metas = load_metadata(meta_path)

    lang_ids, cat_ids, repo_ids = build_ontology(metas)

    spans = build_span_records(
        repos_root=repos_root,
        metas=metas,
        lang_ids=lang_ids,
        cat_ids=cat_ids,
        repo_ids=repo_ids,


main

        "--spans_out",
        type=str,
        required=True,
        help="Output path for spans JSONL file.",
    )
    ap.add_argument(
        "--nodes_out",
        type=str,
        default="",
        help="Optional output path for nodes JSONL (ontology nodes).",
    )
    ap.add_argument(
        "--level_sizes_out",
        type=str,
        default="",
        help="Optional output path for level_sizes JSON file.",
    )
    return ap.parse_args()

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

    repos_root = Path(args.repos_root)
    if not repos_root.is_dir():
        raise FileNotFoundError(f"repos_root does not exist or is not a directory: {repos_root}")

    meta_path = Path(args.metadata_file)
    metas = load_metadata(meta_path)

    lang_ids, cat_ids, repo_ids = build_ontology(metas)

    spans = build_span_records(
        repos_root=repos_root,
        metas=metas,
        lang_ids=lang_ids,
        cat_ids=cat_ids,
        repo_ids=repo_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, lang_ids, cat_ids, repo_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(lang_ids, cat_ids, repo_ids, ls_out)

if __name__ == "__main__":
    main()


"""
Build ArXiv-CLS–style JSONL spans and simple ontology metadata.

This script is a small helper for the ArXiv-CLS variant described in the
TOLBERT paper. It mirrors the style of `build_wos_spans.py` / `build_codehierarchy_spans.py`,
but targets the 2-level arXiv category hierarchy.

The ArXiv-CLS setup in the paper uses:

  - Documents: recent arXiv CS papers (e.g., title + abstract text).
  - Labels: a two-level arXiv category taxonomy, e.g.:
        domain  = "cs"
        category = "cs.LG"

You are expected to provide a metadata file (CSV or JSONL) with at least:

  - doc_id:   unique paper identifier (e.g., arXiv ID)
  - domain:   top-level arXiv area (e.g., "cs", "math")
  - category: fine-grained arXiv category (e.g., "cs.LG", "cs.CL")
  - text:     text span to encode (e.g., title + abstract)

You can override these column names via CLI flags so you can adapt to
your particular preprocessing pipeline.

Outputs
=======

- --spans-out:
    JSONL file with one record per paper of the form:

      {
        "span_id": "arxiv-2101.00001",
        "text": "... title + abstract ...",
        "source_id": "arxiv-2101.00001",
        "node_path": [root_id, domain_id, category_id],
        "meta": {
          "doc_id": "...",
          "domain": "...",
          "category": "..."
        }
      }

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

- --level-sizes-out (optional):
    Small JSON helper:

      {"level_sizes": {1: num_domains, 2: num_categories}}

    which you can paste into a config (see `configs/wos_example.yaml` or
    `configs/codehierarchy_example.yaml` for reference).

The resulting `spans_out` file is directly consumable by
`tolbert.data.TreeOfLifeDataset` and the training skeleton in
`scripts/train_tolbert.py`, as long as your config's `level_sizes`
matches the ontology this script builds.
"""

from __future__ import annotations

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

@dataclass
class ArxivRow:
    doc_id: str
    domain: str
    category: str
    text: str

def _row_to_meta(
    row: Dict[str, object],
    *,
    id_col: str,
    domain_col: str,
    category_col: str,
    text_col: str,
) -> Optional[ArxivRow]:
    try:
        doc_id = str(row[id_col])
        domain = str(row[domain_col])
        category = str(row[category_col])
        text = str(row[text_col])
    except KeyError as exc:
        raise KeyError(f"Missing required column {exc!s} in metadata record: {row}") from exc

    if not text or not text.strip():
        # Skip empty-text rows; they are not useful training examples.
        return None

    return ArxivRow(doc_id=doc_id, domain=domain, category=category, text=text)

def load_arxiv_metadata(
    path: Path,
    *,
    id_col: str,
    domain_col: str,
    category_col: str,
    text_col: str,
) -> List[ArxivRow]:
    """
    Load ArXiv-CLS metadata from CSV or JSON(L).
    """
    rows: List[ArxivRow] = []
    suffix = path.suffix.lower()

    if suffix in {".jsonl", ".json"}:
        with path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                # Allow records nested under "data" for compatibility with some pipelines.
                if isinstance(obj, dict) and "data" in obj and isinstance(obj["data"], dict):
                    obj = obj["data"]
                if not isinstance(obj, dict):
                    continue
                meta = _row_to_meta(
                    obj,
                    id_col=id_col,
                    domain_col=domain_col,
                    category_col=category_col,
                    text_col=text_col,
                )
                if meta is not None:
                    rows.append(meta)
        return rows

    # Default: CSV with header
    with path.open("r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            meta = _row_to_meta(
                row,
                id_col=id_col,
                domain_col=domain_col,
                category_col=category_col,
                text_col=text_col,
            )
            if meta is not None:
                rows.append(meta)
    return rows

def build_ontology(rows: List[ArxivRow]) -> Tuple[Dict[str, int], Dict[str, int]]:
    """
    Build integer node IDs for:

      - level 1:  unique domains (e.g., "cs", "math")
      - level 2:  (domain, category) pairs (e.g., "cs::cs.LG")

    Level conventions:
      level 0: root (id 0)
      level 1: domain nodes
      level 2: category nodes (per domain)
    """
    next_id = 0
    root_id = next_id
    next_id += 1

    domain_ids: Dict[str, int] = {}
    category_ids: Dict[str, int] = {}

    for r in rows:
        if r.domain not in domain_ids:
            domain_ids[r.domain] = next_id
            next_id += 1

        cat_key = f"{r.domain}::{r.category}"
        if cat_key not in category_ids:
            category_ids[cat_key] = next_id
            next_id += 1

    # root_id is currently unused in the returned dicts but documented for clarity
    _ = root_id
    return domain_ids, category_ids

def write_nodes_jsonl(
    out_path: Path,
    domain_ids: Dict[str, int],
    category_ids: Dict[str, int],
) -> None:
    """
    Emit a minimal nodes.jsonl compatible with docs/tree_of_life.md.
    """
    with out_path.open("w", encoding="utf-8") as f:
        # Root (level 0)
        f.write(
            json.dumps(
                {
                    "node_id": 0,
                    "level": 0,
                    "type": "root",
                    "parent_id": None,
                    "name": "Root",
                    "attributes": {},
                }
            )
            + "\n"
        )

        # Domains (level 1)
        for domain_name, nid in domain_ids.items():
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 1,
                        "type": "domain",
                        "parent_id": 0,
                        "name": domain_name,
                        "attributes": {},
                    }
                )
                + "\n"
            )

        # Categories (level 2)
        for key, nid in category_ids.items():
            domain_name, category_name = key.split("::", 1)
            parent_id = domain_ids[domain_name]
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 2,
                        "type": "category",
                        "parent_id": parent_id,
                        "name": category_name,
                        "attributes": {"domain": domain_name},
                    }
                )
                + "\n"
            )

def build_spans(
    rows: List[ArxivRow],
    domain_ids: Dict[str, int],
    category_ids: Dict[str, int],
) -> List[Dict[str, object]]:
    """
    Build span records with node_path = [root_id, domain_id, category_id].
    """
    spans: List[Dict[str, object]] = []
    root_id = 0

    for r in rows:
        domain_id = domain_ids[r.domain]
        cat_key = f"{r.domain}::{r.category}"
        category_id = category_ids[cat_key]

        node_path = [root_id, domain_id, category_id]

        span_id = r.doc_id
        source_id = r.doc_id

        spans.append(
            {
                "span_id": span_id,
                "text": r.text,
                "source_id": source_id,
                "node_path": node_path,
                "meta": {
                    "doc_id": r.doc_id,
                    "domain": r.domain,
                    "category": r.category,
                },
            }
        )

    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(
    domain_ids: Dict[str, int],
    category_ids: Dict[str, int],
    out_path: Path,
) -> None:
    """
    Write level_sizes helper as a dict[int, int] mapping level index
    (excluding the root) to number of classes at that level.

    For ArXiv-CLS we use:
      level 1: domains
      level 2: categories
    """
    level_sizes = {
        1: len(domain_ids),
        2: len(category_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 ArXiv-CLS spans JSONL and ontology metadata.",
    )
    ap.add_argument(
        "--metadata-file",
        type=str,
        required=True,
        help="CSV or JSONL file with at least: doc_id, domain, category, text.",
    )
    ap.add_argument(
        "--spans-out",
        type=str,
        required=True,
        help="Output path for spans JSONL file.",
    )
    ap.add_argument(
        "--nodes-out",
        type=str,
        default="",
        help="Optional output path for nodes JSONL (ontology nodes).",
    )
    ap.add_argument(
        "--level-sizes-out",
        type=str,
        default="",
        help="Optional output path for level_sizes JSON file.",
    )

    # Column name overrides for flexibility.
    ap.add_argument(
        "--id-col",
        type=str,
        default="doc_id",
        help="Column name for document ID.",
    )
    ap.add_argument(
        "--domain-col",
        type=str,
        default="domain",
        help='Column name for top-level arXiv area (e.g., "cs").',
    )
    ap.add_argument(
        "--category-col",
        type=str,
        default="category",
        help='Column name for fine-grained arXiv category (e.g., "cs.LG").',
    )
    ap.add_argument(
        "--text-col",
        type=str,
        default="text",
        help="Column name containing the paper text (e.g., title + abstract).",
    )

    return ap.parse_args()

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

    meta_path = Path(args.metadata_file)
    if not meta_path.is_file():
        raise FileNotFoundError(f"metadata_file does not exist or is not a file: {meta_path}")

    rows = load_arxiv_metadata(
        meta_path,
        id_col=args.id_col,
        domain_col=args.domain_col,
        category_col=args.category_col,
        text_col=args.text_col,
    )
    if not rows:
        raise RuntimeError(f"No valid ArXiv-CLS records were loaded from metadata file: {meta_path}")

    domain_ids, category_ids = build_ontology(rows)

    spans = build_spans(rows, domain_ids=domain_ids, category_ids=category_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, domain_ids, category_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(domain_ids, category_ids, ls_out)

if __name__ == "__main__":
    main()


ArxivRow

    which you can paste into a config (see `configs/wos_example.yaml` or
    `configs/codehierarchy_example.yaml` for reference).

The resulting `spans_out` file is directly consumable by
`tolbert.data.TreeOfLifeDataset` and the training skeleton in
`scripts/train_tolbert.py`, as long as your config's `level_sizes`
matches the ontology this script builds.
"""

from __future__ import annotations

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

@dataclass
class ArxivRow:
    doc_id: str
    domain: str
    category: str
    text: str

def _row_to_meta(
    row: Dict[str, object],
    *,
    id_col: str,
    domain_col: str,
    category_col: str,
    text_col: str,
) -> Optional[ArxivRow]:
    try:
        doc_id = str(row[id_col])
        domain = str(row[domain_col])
        category = str(row[category_col])
        text = str(row[text_col])
    except KeyError as exc:
        raise KeyError(f"Missing required column {exc!s} in metadata record: {row}") from exc

    if not text or not text.strip():
        # Skip empty-text rows; they are not useful training examples.


"""

from __future__ import annotations

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

@dataclass
class ArxivRow:
    doc_id: str
    domain: str
    category: str
    text: str

def _row_to_meta(
    row: Dict[str, object],
    *,
    id_col: str,
    domain_col: str,
    category_col: str,
    text_col: str,
) -> Optional[ArxivRow]:
    try:
        doc_id = str(row[id_col])
        domain = str(row[domain_col])
        category = str(row[category_col])
        text = str(row[text_col])
    except KeyError as exc:
        raise KeyError(f"Missing required column {exc!s} in metadata record: {row}") from exc

    if not text or not text.strip():
        # Skip empty-text rows; they are not useful training examples.
        return None

    return ArxivRow(doc_id=doc_id, domain=domain, category=category, text=text)

def load_arxiv_metadata(
    path: Path,
    *,
    id_col: str,
    domain_col: str,
    category_col: str,
    text_col: str,
) -> List[ArxivRow]:
    """
    Load ArXiv-CLS metadata from CSV or JSON(L).
    """
    rows: List[ArxivRow] = []
    suffix = path.suffix.lower()

    if suffix in {".jsonl", ".json"}:
        with path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()


    id_col: str,
    domain_col: str,
    category_col: str,
    text_col: str,
) -> Optional[ArxivRow]:
    try:
        doc_id = str(row[id_col])
        domain = str(row[domain_col])
        category = str(row[category_col])
        text = str(row[text_col])
    except KeyError as exc:
        raise KeyError(f"Missing required column {exc!s} in metadata record: {row}") from exc

    if not text or not text.strip():
        # Skip empty-text rows; they are not useful training examples.
        return None

    return ArxivRow(doc_id=doc_id, domain=domain, category=category, text=text)

def load_arxiv_metadata(
    path: Path,
    *,
    id_col: str,
    domain_col: str,
    category_col: str,
    text_col: str,
) -> List[ArxivRow]:
    """
    Load ArXiv-CLS metadata from CSV or JSON(L).
    """
    rows: List[ArxivRow] = []
    suffix = path.suffix.lower()

    if suffix in {".jsonl", ".json"}:
        with path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                # Allow records nested under "data" for compatibility with some pipelines.
                if isinstance(obj, dict) and "data" in obj and isinstance(obj["data"], dict):
                    obj = obj["data"]
                if not isinstance(obj, dict):
                    continue
                meta = _row_to_meta(
                    obj,
                    id_col=id_col,
                    domain_col=domain_col,
                    category_col=category_col,
                    text_col=text_col,
                )
                if meta is not None:
                    rows.append(meta)
        return rows

    # Default: CSV with header
    with path.open("r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            meta = _row_to_meta(
                row,
                id_col=id_col,
                domain_col=domain_col,
                category_col=category_col,
                text_col=text_col,
            )
            if meta is not None:
                rows.append(meta)
    return rows

def build_ontology(rows: List[ArxivRow]) -> Tuple[Dict[str, int], Dict[str, int]]:
    """
    Build integer node IDs for:

      - level 1:  unique domains (e.g., "cs", "math")
      - level 2:  (domain, category) pairs (e.g., "cs::cs.LG")

    Level conventions:
      level 0: root (id 0)
      level 1: domain nodes
      level 2: category nodes (per domain)
    """
    next_id = 0
    root_id = next_id
    next_id += 1

    domain_ids: Dict[str, int] = {}
    category_ids: Dict[str, int] = {}


                if meta is not None:
                    rows.append(meta)
        return rows

    # Default: CSV with header
    with path.open("r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            meta = _row_to_meta(
                row,
                id_col=id_col,
                domain_col=domain_col,
                category_col=category_col,
                text_col=text_col,
            )
            if meta is not None:
                rows.append(meta)
    return rows

def build_ontology(rows: List[ArxivRow]) -> Tuple[Dict[str, int], Dict[str, int]]:
    """
    Build integer node IDs for:

      - level 1:  unique domains (e.g., "cs", "math")
      - level 2:  (domain, category) pairs (e.g., "cs::cs.LG")

    Level conventions:
      level 0: root (id 0)
      level 1: domain nodes
      level 2: category nodes (per domain)
    """
    next_id = 0
    root_id = next_id
    next_id += 1

    domain_ids: Dict[str, int] = {}
    category_ids: Dict[str, int] = {}

    for r in rows:
        if r.domain not in domain_ids:
            domain_ids[r.domain] = next_id
            next_id += 1

        cat_key = f"{r.domain}::{r.category}"
        if cat_key not in category_ids:
            category_ids[cat_key] = next_id
            next_id += 1

    # root_id is currently unused in the returned dicts but documented for clarity
    _ = root_id
    return domain_ids, category_ids

def write_nodes_jsonl(
    out_path: Path,
    domain_ids: Dict[str, int],
    category_ids: Dict[str, int],
) -> None:
    """
    Emit a minimal nodes.jsonl compatible with docs/tree_of_life.md.
    """
    with out_path.open("w", encoding="utf-8") as f:
        # Root (level 0)
        f.write(
            json.dumps(
                {
                    "node_id": 0,
                    "level": 0,
                    "type": "root",
                    "parent_id": None,
                    "name": "Root",


    next_id += 1

    domain_ids: Dict[str, int] = {}
    category_ids: Dict[str, int] = {}

    for r in rows:
        if r.domain not in domain_ids:
            domain_ids[r.domain] = next_id
            next_id += 1

        cat_key = f"{r.domain}::{r.category}"
        if cat_key not in category_ids:
            category_ids[cat_key] = next_id
            next_id += 1

    # root_id is currently unused in the returned dicts but documented for clarity
    _ = root_id
    return domain_ids, category_ids

def write_nodes_jsonl(
    out_path: Path,
    domain_ids: Dict[str, int],
    category_ids: Dict[str, int],
) -> None:
    """
    Emit a minimal nodes.jsonl compatible with docs/tree_of_life.md.
    """
    with out_path.open("w", encoding="utf-8") as f:
        # Root (level 0)
        f.write(
            json.dumps(
                {
                    "node_id": 0,
                    "level": 0,
                    "type": "root",
                    "parent_id": None,
                    "name": "Root",
                    "attributes": {},
                }
            )
            + "\n"
        )

        # Domains (level 1)
        for domain_name, nid in domain_ids.items():
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 1,
                        "type": "domain",
                        "parent_id": 0,
                        "name": domain_name,
                        "attributes": {},
                    }
                )
                + "\n"
            )

        # Categories (level 2)
        for key, nid in category_ids.items():
            domain_name, category_name = key.split("::", 1)
            parent_id = domain_ids[domain_name]
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 2,
                        "type": "category",
                        "parent_id": parent_id,
                        "name": category_name,
                        "attributes": {"domain": domain_name},
                    }
                )
                + "\n"
            )

def build_spans(
    rows: List[ArxivRow],
    domain_ids: Dict[str, int],
    category_ids: Dict[str, int],
) -> List[Dict[str, object]]:
    """
    Build span records with node_path = [root_id, domain_id, category_id].
    """
    spans: List[Dict[str, object]] = []
    root_id = 0

    for r in rows:
        domain_id = domain_ids[r.domain]
        cat_key = f"{r.domain}::{r.category}"
        category_id = category_ids[cat_key]

        node_path = [root_id, domain_id, category_id]


        # Categories (level 2)
        for key, nid in category_ids.items():
            domain_name, category_name = key.split("::", 1)
            parent_id = domain_ids[domain_name]
            f.write(
                json.dumps(
                    {
                        "node_id": nid,
                        "level": 2,
                        "type": "category",
                        "parent_id": parent_id,
                        "name": category_name,
                        "attributes": {"domain": domain_name},
                    }
                )
                + "\n"
            )

def build_spans(
    rows: List[ArxivRow],
    domain_ids: Dict[str, int],
    category_ids: Dict[str, int],
) -> List[Dict[str, object]]:
    """
    Build span records with node_path = [root_id, domain_id, category_id].
    """
    spans: List[Dict[str, object]] = []
    root_id = 0

    for r in rows:
        domain_id = domain_ids[r.domain]
        cat_key = f"{r.domain}::{r.category}"
        category_id = category_ids[cat_key]

        node_path = [root_id, domain_id, category_id]

        span_id = r.doc_id
        source_id = r.doc_id

        spans.append(
            {
                "span_id": span_id,
                "text": r.text,
                "source_id": source_id,
                "node_path": node_path,
                "meta": {
                    "doc_id": r.doc_id,
                    "domain": r.domain,
                    "category": r.category,
                },
            }
        )

    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(
    domain_ids: Dict[str, int],
    category_ids: Dict[str, int],
    out_path: Path,
) -> None:
    """
    Write level_sizes helper as a dict[int, int] mapping level index
    (excluding the root) to number of classes at that level.

    For ArXiv-CLS we use:
      level 1: domains
      level 2: categories


        span_id = r.doc_id
        source_id = r.doc_id

        spans.append(
            {
                "span_id": span_id,
                "text": r.text,
                "source_id": source_id,
                "node_path": node_path,
                "meta": {
                    "doc_id": r.doc_id,
                    "domain": r.domain,
                    "category": r.category,
                },
            }
        )

    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(
    domain_ids: Dict[str, int],
    category_ids: Dict[str, int],
    out_path: Path,
) -> None:
    """
    Write level_sizes helper as a dict[int, int] mapping level index
    (excluding the root) to number of classes at that level.

    For ArXiv-CLS we use:
      level 1: domains
      level 2: categories
    """
    level_sizes = {
        1: len(domain_ids),
        2: len(category_ids),
    }
    out = {"level_sizes": level_sizes}


                "text": r.text,
                "source_id": source_id,
                "node_path": node_path,
                "meta": {
                    "doc_id": r.doc_id,
                    "domain": r.domain,
                    "category": r.category,
                },
            }
        )

    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(
    domain_ids: Dict[str, int],
    category_ids: Dict[str, int],
    out_path: Path,
) -> None:
    """
    Write level_sizes helper as a dict[int, int] mapping level index
    (excluding the root) to number of classes at that level.

    For ArXiv-CLS we use:
      level 1: domains
      level 2: categories
    """
    level_sizes = {
        1: len(domain_ids),
        2: len(category_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 ArXiv-CLS spans JSONL and ontology metadata.",
    )
    ap.add_argument(
        "--metadata-file",
        type=str,
        required=True,
        help="CSV or JSONL file with at least: doc_id, domain, category, text.",
    )
    ap.add_argument(
        "--spans-out",
        type=str,
        required=True,
        help="Output path for spans JSONL file.",
    )
    ap.add_argument(
        "--nodes-out",


    category_ids: Dict[str, int],
    out_path: Path,
) -> None:
    """
    Write level_sizes helper as a dict[int, int] mapping level index
    (excluding the root) to number of classes at that level.

    For ArXiv-CLS we use:
      level 1: domains
      level 2: categories
    """
    level_sizes = {
        1: len(domain_ids),
        2: len(category_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 ArXiv-CLS spans JSONL and ontology metadata.",
    )
    ap.add_argument(
        "--metadata-file",
        type=str,
        required=True,
        help="CSV or JSONL file with at least: doc_id, domain, category, text.",
    )
    ap.add_argument(
        "--spans-out",
        type=str,
        required=True,
        help="Output path for spans JSONL file.",
    )
    ap.add_argument(
        "--nodes-out",
        type=str,
        default="",
        help="Optional output path for nodes JSONL (ontology nodes).",
    )
    ap.add_argument(
        "--level-sizes-out",
        type=str,
        default="",
        help="Optional output path for level_sizes JSON file.",
    )

    # Column name overrides for flexibility.
    ap.add_argument(
        "--id-col",
        type=str,
        default="doc_id",
        help="Column name for document ID.",
    )
    ap.add_argument(
        "--domain-col",
        type=str,
        default="domain",
        help='Column name for top-level arXiv area (e.g., "cs").',
    )
    ap.add_argument(
        "--category-col",
        type=str,
        default="category",
        help='Column name for fine-grained arXiv category (e.g., "cs.LG").',
    )
    ap.add_argument(
        "--text-col",
        type=str,
        default="text",
        help="Column name containing the paper text (e.g., title + abstract).",
    )

    return ap.parse_args()

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

    meta_path = Path(args.metadata_file)
    if not meta_path.is_file():
        raise FileNotFoundError(f"metadata_file does not exist or is not a file: {meta_path}")

    rows = load_arxiv_metadata(
        meta_path,
        id_col=args.id_col,
        domain_col=args.domain_col,
        category_col=args.category_col,
        text_col=args.text_col,
    )
    if not rows:
        raise RuntimeError(f"No valid ArXiv-CLS records were loaded from metadata file: {meta_path}")

    domain_ids, category_ids = build_ontology(rows)


main

        type=str,
        default="domain",
        help='Column name for top-level arXiv area (e.g., "cs").',
    )
    ap.add_argument(
        "--category-col",
        type=str,
        default="category",
        help='Column name for fine-grained arXiv category (e.g., "cs.LG").',
    )
    ap.add_argument(
        "--text-col",
        type=str,
        default="text",
        help="Column name containing the paper text (e.g., title + abstract).",
    )

    return ap.parse_args()

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

    meta_path = Path(args.metadata_file)
    if not meta_path.is_file():
        raise FileNotFoundError(f"metadata_file does not exist or is not a file: {meta_path}")

    rows = load_arxiv_metadata(
        meta_path,
        id_col=args.id_col,
        domain_col=args.domain_col,
        category_col=args.category_col,
        text_col=args.text_col,
    )
    if not rows:
        raise RuntimeError(f"No valid ArXiv-CLS records were loaded from metadata file: {meta_path}")

    domain_ids, category_ids = build_ontology(rows)

    spans = build_spans(rows, domain_ids=domain_ids, category_ids=category_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, domain_ids, category_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(domain_ids, category_ids, ls_out)

if __name__ == "__main__":
    main()