"""Custom physical-support track: parent activity from noisy nearly-colliding rays."""
import numpy as np

META = {
    "name": "collision_support",
    "domain": "embedding",
    "description": "Parent-block activity from replicated noisy observations of two nearly colliding sign-invariant child rays.",
}


def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(seed)
    s = 0.20
    sigma = 0.35

    def make(n):
        active = rng.integers(0, 2, size=n).astype(np.int64)
        x = rng.normal(0, sigma, size=(n, 4, 2)).astype(np.float32)
        a = np.arcsin(s / 2.0)
        rays = np.array([[np.cos(a), np.sin(a)], [np.cos(a), -np.sin(a)]], dtype=np.float32)
        for i in range(n):
            if active[i]:
                for t in range(4):
                    k = rng.integers(0, 2)
                    sign = rng.choice([-1.0, 1.0])
                    x[i, t] += sign * rays[k]
        return x.reshape(n, -1), active

    xtr, ytr = make(n_train)
    xte, yte = make(n_test)
    return {
        "xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
        "task": "classification", "metric": "err",
        "input_shape": (8,), "out_dim": 2,
        "s": s, "sigma": sigma, "N": n_train, "T": 4,
    }