"""Local stage-2 track for branch-length-aware rooted-tree encoders."""
import numpy as np

META = {
    "name": "rooted_tree_branch_classification",
    "domain": "graph-nn",
    "description": "Classification of rooted trees with identical edge-length multisets but different child topology; inputs encode parent indices and positive branch lengths.",
}

# Six vertices: root 0 and five descendants. Both classes have five edges.
PARENTS = {
    0: np.array([-1, 0, 0, 0, 3, 3]),  # root has 3 children; node 3 has 2
    1: np.array([-1, 0, 1, 1, 1, 1]),  # root has 1 child; node 1 has 4
}

def _make(rng, n):
    xs, ys = [], []
    base = np.array([0.45, 0.72, 1.05, 1.38, 1.76], dtype=np.float32)
    for _ in range(int(n)):
        label = int(rng.randint(0, 2))
        vals = rng.permutation(base)
        # Small common nuisance scale preserves equal multiset within each sample.
        vals = vals * rng.uniform(0.92, 1.08)
        parent = PARENTS[label]
        x = np.zeros((6, 2), dtype=np.float32)
        x[:, 0] = np.where(parent < 0, -1.0, parent / 5.0)
        edge_nodes = np.flatnonzero(parent >= 0)
        x[edge_nodes, 1] = vals
        xs.append(x); ys.append(label)
    return np.asarray(xs), np.asarray(ys, dtype=np.int64)

def get_dataset(seed, n_train, n_test):
    xtr, ytr = _make(np.random.RandomState(int(seed)), n_train)
    xte, yte = _make(np.random.RandomState(int(seed) + 100003), n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "classification", "metric": "err", "out_dim": 2}