src/synthetic_tree.py

"""
Synthetic random tree, for validating the embedding/training/evaluation
pipeline WITHOUT needing live PBDB network access (blocked in this
development sandbox). Provenance is explicitly "SYNTHETIC_TREE" — never
to be confused with real PBDB taxonomy data, and never silently
substituted for it (same discipline as get_synthetic_dataset() in
data_real.py: explicit opt-in only, clearly labeled).

Construction: a random recursive tree — node i (for i >= 1) attaches to
a uniformly random existing node in {0, ..., i-1}. This guarantees a
single connected tree with no cycles, node 0 as root, and is a standard,
well-understood random tree model (not an arbitrary ad hoc generator).
"""
from __future__ import annotations
from typing import Dict, List, Tuple

import torch

from .data_pbdb_taxonomy import TaxonomyEdgeDataset


def generate_synthetic_tree(
    n_nodes: int = 300,
    seed: int = 0,
) -> Tuple[List[Tuple[str, str]], Dict[str, Dict[str, int]]]:
    gen = torch.Generator().manual_seed(seed)
    edges: List[Tuple[str, str]] = []
    node_attrs: Dict[str, Dict[str, int]] = {"n0": {"depth": 0}}
    depth = [0] * n_nodes
    for i in range(1, n_nodes):
        parent = int(torch.randint(0, i, (1,), generator=gen).item())
        edges.append((f"n{i}", f"n{parent}"))
        depth[i] = depth[parent] + 1
        node_attrs[f"n{i}"] = {"depth": depth[i]}
    return edges, node_attrs


def generate_balanced_tree(
    branching_factor: int = 3,
    depth: int = 6,
) -> Tuple[List[Tuple[str, str]], Dict[str, Dict[str, int]]]:

    edges: List[Tuple[str, str]] = []
    node_attrs: Dict[str, Dict[str, int]] = {"n0": {"depth": 0}}
    frontier = ["n0"]
    counter = 1
    for d in range(1, depth + 1):
        next_frontier = []
        for parent in frontier:
            for _ in range(branching_factor):
                name = f"n{counter}"
                counter += 1
                edges.append((name, parent))
                node_attrs[name] = {"depth": d}
                next_frontier.append(name)
        frontier = next_frontier
    return edges, node_attrs


def get_synthetic_tree_dataset(n_nodes: int = 300, seed: int = 0, tree_type: str = "random_recursive",
                                 branching_factor: int = 3, depth: int = 6):
    if tree_type == "balanced":
        edges, node_attrs = generate_balanced_tree(branching_factor=branching_factor, depth=depth)
    else:
        edges, node_attrs = generate_synthetic_tree(n_nodes=n_nodes, seed=seed)
    dataset = TaxonomyEdgeDataset(edges, node_attrs)
    node_depths_by_idx = {
        dataset.node_to_idx[name]: attrs["depth"] for name, attrs in node_attrs.items()
    }
    meta = {
        "num_nodes": dataset.num_nodes,
        "num_edges": len(edges),
        "seed": seed,
        "tree_type": tree_type,
        "node_depths_by_idx": node_depths_by_idx,
    }
    print(f"[data] SYNTHETIC_TREE ({tree_type}): {dataset.num_nodes} nodes, {len(edges)} edges "
          f"(seed={seed}) — NOT real PBDB data")
    return dataset, meta, "SYNTHETIC_TREE"