src/data_pbdb_taxonomy.py

"""
PBDB taxonomic hierarchy, built from the SAME verified endpoint as
data_pbdb.py (occs/list with show=classext), rather than the separate
`/data1.2/taxa/list` endpoint an earlier draft plan sketched — that
endpoint's field names (oid/nam/rnk/par) were explicitly marked "example;
finalize after --discover" in that sketch, i.e. unverified even by its
own author. Reusing occs/list is lower-risk: its schema (for this
purpose) is now backed by an actual live-data example found this session
(see data_pbdb.py's module docstring) rather than a guess.

CONFIRMED this session (real example row from the paleobioDB R package
docs, not just prose): show=classext returns `phylum`, `class`, `order`,
`family`, `genus` as direct columns on each occurrence record. Standard
Linnaean rank order (phylum > class > order > family > genus) is used to
derive parent-child edges from each occurrence's classification path.

This is real taxonomic signal (whatever PBDB's curators have assigned),
not synthesized — but it is also incomplete in a specific, honest way:
it only contains ranks that co-occur on real fossil occurrence records,
so higher clades with no occurrences at a given rank won't appear, and a
genus is linked to whichever family/order/class/phylum it was reported
under (which is occasionally inconsistent between occurrences for
taxonomically disputed groups — this module does not attempt to resolve
those disputes, it records them as-observed and lets a human decide, per
the project's no-fabrication rule).
"""
from __future__ import annotations
import hashlib
import json
from typing import Any, Dict, List, Optional, Sequence, Tuple

from torch.utils.data import Dataset

from .data_pbdb import fetch_occurrences
from .provenance import DataLoadError, SchemaValidationError

RANK_CHAIN = ("phylum", "class", "order", "family", "genus")

TAXONOMY_REQUIRED_FIELDS = ("occurrence_no",)


def discover_taxonomy_schema(base_name: str = "Canidae", limit: int = 5) -> Dict[str, Any]:
    from .data_pbdb import discover_schema
    result = discover_schema(base_name=base_name, limit=limit, show="classext")
    observed = set(result["observed_fields"])
    result["rank_fields_present"] = [f for f in RANK_CHAIN if f in observed]
    result["rank_fields_missing"] = [f for f in RANK_CHAIN if f not in observed]
    return result


def build_edge_list(
    records: List[Dict[str, Any]],
    min_rank: Optional[str] = None,
) -> Tuple[List[Tuple[str, str]], Dict[str, Dict[str, Any]]]:
    chain = RANK_CHAIN
    if min_rank is not None:
        if min_rank not in RANK_CHAIN:
            raise ValueError(f"min_rank must be one of {RANK_CHAIN}, got {min_rank!r}")
        chain = RANK_CHAIN[: RANK_CHAIN.index(min_rank) + 1]

    edges = set()
    node_attrs: Dict[str, Dict[str, Any]] = {}
    for r in records:
        path = [r.get(rank) for rank in chain]
        # Walk consecutive (parent_rank, child_rank) pairs in this record's
        # path; only add an edge where both ends are real, non-empty
        # strings taken from this SAME record (never mixed across records).
        for i in range(len(path) - 1):
            parent, child = path[i], path[i + 1]
            if not parent or not child:
                continue
            if not isinstance(parent, str) or not isinstance(child, str):
                continue
            edges.add((child, parent))
            node_attrs[parent] = {"rank": chain[i]}
            node_attrs[child] = {"rank": chain[i + 1]}

    edge_list = sorted(edges)
    if len(edge_list) < 2:
        raise SchemaValidationError(
            f"only {len(edge_list)} usable (child, parent) edge(s) could be "
            f"built from {len(records)} occurrence records — not enough "
            f"classification coverage to build a meaningful tree. Try a "
            f"broader base_name or check discover_taxonomy_schema() output.",
            outcome_code="INSUFFICIENT_TAXONOMY_EDGES",
        )
    return edge_list, node_attrs


def hash_edge_list(edges: List[Tuple[str, str]]) -> str:
    blob = json.dumps(sorted(edges), sort_keys=True).encode()
    return hashlib.sha256(blob).hexdigest()


class TaxonomyEdgeDataset(Dataset):

    def __init__(self, edges: List[Tuple[str, str]], node_attrs: Dict[str, Dict[str, Any]]):
        self.edges = edges
        self.node_attrs = node_attrs
        nodes = sorted(node_attrs.keys())
        self.node_to_idx = {n: i for i, n in enumerate(nodes)}
        self.idx_to_node = {i: n for n, i in self.node_to_idx.items()}
        self.num_nodes = len(nodes)
        self.edge_idx = [(self.node_to_idx[c], self.node_to_idx[p]) for c, p in edges]

    def __len__(self):
        return len(self.edge_idx)

    def __getitem__(self, idx):
        child_idx, parent_idx = self.edge_idx[idx]
        return {"child_idx": child_idx, "parent_idx": parent_idx}


def get_pbdb_taxonomy_dataset(
    base_names: Sequence[str] = ("Dinosauria", "Mammalia"),
    max_taxa_per_group: int = 8000,
    min_rank: Optional[str] = None,
    cache_dir: Optional[str] = "data/pbdb_taxonomy_cache",
):
    all_records = []
    for base_name in base_names:
        recs = fetch_occurrences(
            base_name=base_name, max_records=max_taxa_per_group,
            show="classext", cache_dir=cache_dir,
            required_fields=TAXONOMY_REQUIRED_FIELDS,
        )
        all_records.extend(recs)

    edges, node_attrs = build_edge_list(all_records, min_rank=min_rank)
    dataset = TaxonomyEdgeDataset(edges, node_attrs)
    edge_hash = hash_edge_list(edges)
    meta = {
        "num_nodes": dataset.num_nodes,
        "num_edges": len(edges),
        "base_names": list(base_names),
        "edge_hash": edge_hash,
        "node_to_idx": dataset.node_to_idx,
    }
    print(f"[data] REAL_PBDB_TAXONOMY: {dataset.num_nodes} nodes, "
          f"{len(edges)} edges from {len(all_records)} occurrence records "
          f"across {list(base_names)}")
    return dataset, meta, "REAL_PBDB_TAXONOMY"