src/run_hierarchy_embed.py

"""
Track B driver: train Poincaré vs Euclidean hierarchy embeddings under
identical conditions and report the comparison table.

    python -m src.run_hierarchy_embed --synthetic --n-nodes 500 --dim 8
    python -m src.run_hierarchy_embed --pbdb-taxa Dinosauria Mammalia --dim 8
"""
from __future__ import annotations
import argparse
import random
from typing import List, Tuple

import numpy as np
import torch

from .embed_hierarchy import train_hierarchy_embedding
from .eval_hierarchy import compare_geometries, radius_diagnostics
from .synthetic_tree import get_synthetic_tree_dataset
from .data_pbdb_taxonomy import get_pbdb_taxonomy_dataset
from .provenance import CheckpointStore, hash_code
import os


def set_all_seeds(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def split_edges(dataset, test_frac: float = 0.2, seed: int = 0) -> Tuple[List, List]:
    """
    WARNING, found by actually running this on a real tree (not assumed):
    for a TREE, every non-root node has EXACTLY ONE parent edge (verified:
    max count of any node as a 'child' across all edges is 1). Holding out
    that single edge removes 100% of that node's positive training signal
    -- it has nothing left pulling it toward its true parent. This makes a
    held-out-edge split fundamentally unsuited to evaluating a tree
    embedding's generalization; it doesn't test what you think it tests.
    Use eval_mode='reconstruction' (the default) instead, which matches
    Nickel & Kiela's own protocol: train on all edges, evaluate rank
    recovery on those same edges (measuring embedding CAPACITY/fidelity at
    a given dimension, not generalization to unseen relations -- which is
    what "does hyperbolic geometry need fewer dimensions" actually means).
    This function is kept only for the (explicitly discouraged) held_out
    eval mode.
    """
    edges = list(dataset.edge_idx)
    rng = random.Random(seed)
    rng.shuffle(edges)
    n_test = max(1, int(len(edges) * test_frac))
    return edges[n_test:], edges[:n_test]


def main():
    p = argparse.ArgumentParser()
    src = p.add_mutually_exclusive_group(required=True)
    src.add_argument("--synthetic", action="store_true",
                      help="Use a synthetic random tree (runs anywhere, no network).")
    src.add_argument("--pbdb-taxa", nargs="+", default=None,
                      help="Real PBDB taxon groups (requires network to paleobiodb.org).")
    p.add_argument("--tree-type", choices=["random_recursive", "balanced"], default="random_recursive",
                    help="'random_recursive': shallow/bushy, depth~ln(n). "
                         "'balanced': fixed branching_factor/depth, genuinely deep -- "
                         "the regime the hyperbolic-advantage literature targets.")
    p.add_argument("--branching-factor", type=int, default=3)
    p.add_argument("--tree-depth", type=int, default=6)
    p.add_argument("--n-nodes", type=int, default=500, help="Synthetic tree size (random_recursive only).")
    p.add_argument("--dims", type=int, nargs="+", default=[8],
                    help="One or more embedding dimensions to sweep.")
    p.add_argument("--epochs", type=int, default=200)
    p.add_argument("--lr", type=float, default=0.01)
    p.add_argument("--margin", type=float, default=1.0)
    p.add_argument("--neg-samples", type=int, default=10)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--seeds", type=int, nargs="+", default=None,
                    help="If given, run all (dim, loss_type) combos across these "
                         "seeds and report mean +/- std, instead of a single seed.")
    p.add_argument("--loss-types", choices=["margin", "softmax"], nargs="+", default=["margin"],
                    help="One or both losses to compare under identical conditions.")
    p.add_argument("--burn-in-epochs", type=int, default=0)
    p.add_argument("--burn-in-lr-mult", type=float, default=0.1)
    p.add_argument("--learnable-c", action="store_true",
                    help="Let curvature be learned per model instead of fixed at --c. "
                         "Applies to the Poincare model only (a fixed scalar for the "
                         "Euclidean baseline has no equivalent meaning). Guarded by "
                         "--c-min/--c-max and a separate curvature learning rate -- "
                         "see embed_hierarchy.py for why those guards are required.")
    p.add_argument("--c", type=float, default=1.0, help="Fixed curvature when --learnable-c is not set.")
    p.add_argument("--c-values", type=float, nargs="+", default=None,
                    help="Sweep a fixed-c grid instead of a single --c (recommended "
                         "primary methodology: interpretable, no coupled instability, "
                         "matches how most hyperbolic embedding papers report results). "
                         "Mutually exclusive with --learnable-c.")
    p.add_argument("--c-min", type=float, default=0.1)
    p.add_argument("--c-max", type=float, default=3.0)
    p.add_argument("--curvature-lr-mult", type=float, default=0.1)
    p.add_argument("--optimizer", choices=["radam", "rsgd"], default="radam",
                    help="radam (RiemannianAdam, used throughout so far) or "
                         "rsgd (RiemannianSGD -- what Nickel & Kiela's original "
                         "2017 paper actually used; never isolated as a variable "
                         "in this project before now).")
    p.add_argument("--eval-mode", choices=["reconstruction", "held_out"], default="reconstruction",
                    help="'reconstruction' (default, matches Nickel & Kiela's protocol): "
                         "train on all edges, evaluate rank-recovery on those same edges "
                         "-- measures embedding capacity/fidelity at a given dimension. "
                         "'held_out': held-out edge split -- WARNING, verified broken for "
                         "tree data (every node has exactly one parent edge, so holding "
                         "it out removes 100%% of that node's training signal). Kept only "
                         "for illustration of that failure mode.")
    args = p.parse_args()

    set_all_seeds(args.seed)

    if args.synthetic:
        dataset, meta, provenance = get_synthetic_tree_dataset(
            n_nodes=args.n_nodes, seed=args.seed, tree_type=args.tree_type,
            branching_factor=args.branching_factor, depth=args.tree_depth,
        )
        node_depths = meta["node_depths_by_idx"]
    else:
        dataset, meta, provenance = get_pbdb_taxonomy_dataset(base_names=args.pbdb_taxa)
        node_depths = None  # PBDB tree has no single global root; depth diagnostics skipped

    if args.eval_mode == "reconstruction":
        train_edges = list(dataset.edge_idx)
        test_edges = train_edges  # by design: evaluating recovery of what was trained on
    else:
        print("\n*** WARNING: --eval-mode held_out is known-broken for tree data. ***")
        print("*** Every node has exactly one parent edge; holding it out removes  ***")
        print("*** 100% of that node's training signal. Numbers below will look    ***")
        print("*** close to random and should NOT be read as 'geometry doesn't help'.***\n")
        train_edges, test_edges = split_edges(dataset, test_frac=0.2, seed=args.seed)
    print(f"[data] provenance={provenance} nodes={dataset.num_nodes} "
          f"train_edges={len(train_edges)} test_edges={len(test_edges)}")

    class _EdgeSubset:
        def __init__(self, edge_idx, num_nodes):
            self.edge_idx = edge_idx
            self.num_nodes = num_nodes
    train_ds = _EdgeSubset(train_edges, dataset.num_nodes)

    if args.learnable_c and args.c_values:
        p.error("--learnable-c and --c-values are mutually exclusive")

    results_table = []
    seeds = args.seeds or [args.seed]
    c_grid = args.c_values or [args.c]  # single value unless a grid was given
    for dim in args.dims:
        for loss_type in args.loss_types:
            for c_val in c_grid:
                label = f"dim={dim}, loss={loss_type}, " + (
                    "learnable_c" if args.learnable_c else f"c={c_val}"
                )
                print(f"\n{'='*64}\n{label}\n{'='*64}")
                deltas_mrr, deltas_rank = [], []
                per_seed = []
                for sd in seeds:
                    p_model, p_metrics = train_hierarchy_embedding(
                        train_ds, geometry="poincare", dim=dim, epochs=args.epochs,
                        lr=args.lr, margin=args.margin, neg_samples=args.neg_samples,
                        loss_type=loss_type, burn_in_epochs=args.burn_in_epochs,
                        burn_in_lr_mult=args.burn_in_lr_mult, c=c_val,
                        learnable_c=args.learnable_c, curvature_lr_mult=args.curvature_lr_mult,
                        optimizer_type=args.optimizer,
                        c_min=args.c_min, c_max=args.c_max, seed=sd,
                    )
                    e_model, e_metrics = train_hierarchy_embedding(
                        train_ds, geometry="euclidean", dim=dim, epochs=args.epochs,
                        lr=args.lr, margin=args.margin, neg_samples=args.neg_samples,
                        loss_type=loss_type, burn_in_epochs=args.burn_in_epochs,
                        burn_in_lr_mult=args.burn_in_lr_mult, seed=sd,
                    )
                    comparison = compare_geometries(p_model, e_model, test_edges)
                    deltas_mrr.append(comparison["delta_mrr"])
                    deltas_rank.append(comparison["delta_mean_rank"])
                    per_seed.append(comparison)
                    print(f"  seed={sd:<5} P_mrr={comparison['poincare']['mrr']:.4f} "
                          f"E_mrr={comparison['euclidean']['mrr']:.4f} "
                          f"delta_mrr={comparison['delta_mrr']:+.4f}")

                deltas_mrr_t = torch.tensor(deltas_mrr)
                mean_delta = deltas_mrr_t.mean().item()
                std_delta = deltas_mrr_t.std().item() if len(seeds) > 1 else 0.0
                wins = sum(1 for d in deltas_mrr if d > 0)
                print(f"  --> mean delta_mrr = {mean_delta:+.4f} +/- {std_delta:.4f} "
                      f"(Poincare wins {wins}/{len(seeds)} seeds)")

                if node_depths is not None:
                    p_radius = radius_diagnostics(p_model, node_depths)
                    e_radius = radius_diagnostics(e_model, node_depths)
                    print(f"  radius-depth correlation  poincare={p_radius['radius_depth_correlation']:+.4f}"
                          f"  euclidean={e_radius['radius_depth_correlation']:+.4f}")

                if args.learnable_c:
                    print(f"  learned curvature (last seed): {p_model.manifold.c.item():.4f} "
                          f"(started at {c_val})")

                results_table.append({
                    "dim": dim, "loss_type": loss_type, "c": c_val,
                    "learnable_c": args.learnable_c, "seeds": seeds,
                    "mean_delta_mrr": mean_delta, "std_delta_mrr": std_delta,
                    "wins": wins, "n_seeds": len(seeds),
                })

    print(f"\n{'='*64}\nSummary across dimensions, losses, and curvatures\n{'='*64}")
    print(f"{'dim':>5} {'loss':>10} {'c':>8} {'mean_delta_mrr':>16} {'std':>8} {'wins':>8}")
    for r in results_table:
        c_label = "learned" if r["learnable_c"] else f"{r['c']:.3f}"
        print(f"{r['dim']:>5} {r['loss_type']:>10} {c_label:>8} {r['mean_delta_mrr']:>+16.4f} "
              f"{r['std_delta_mrr']:>8.4f} {r['wins']}/{r['n_seeds']:>6}")

    store = CheckpointStore(checkpoints_dir="checkpoints")
    code_hash = hash_code(os.path.dirname(os.path.abspath(__file__)))
    store.save(
        model_state={"results_table": results_table},
        config={"dims": args.dims, "epochs": args.epochs, "lr": args.lr,
                "margin": args.margin, "neg_samples": args.neg_samples, "seed": args.seed},
        dataset_hash=meta.get("edge_hash") or f"synthetic_{meta.get('seed')}_{dataset.num_nodes}",
        code_hash=code_hash,
        data_provenance=provenance,
        extra={"meta": {k: v for k, v in meta.items() if k != "node_to_idx"}},
    )


if __name__ == "__main__":
    main()