import numpy as np

META = {
    "name": "jacobian_simplex_triangle",
    "domain": "geometric_graph",
    "description": "Invariant triangle edge lengths with Heron area and angle-sensitive regression target.",
}

def get_dataset(seed, n_train, n_test):
    def sample(n, s):
        rng = np.random.RandomState(s)
        p = rng.normal(size=(n, 3, 3)).astype(np.float32)
        near = rng.rand(n) < 0.25
        if near.any():
            p[near, 2] = p[near, 1] + 0.04 * rng.normal(size=(near.sum(), 3)).astype(np.float32)
        p *= (0.6 + 1.4 * rng.rand(n, 1, 1)).astype(np.float32)
        a = np.linalg.norm(p[:, 0] - p[:, 1], axis=1)
        b = np.linalg.norm(p[:, 1] - p[:, 2], axis=1)
        c = np.linalg.norm(p[:, 2] - p[:, 0], axis=1)
        sides = np.stack([a, b, c], axis=1).astype(np.float32)
        ss = sides.sum(1) / 2
        rad = ss * (ss-a) * (ss-b) * (ss-c)
        area = np.sqrt(np.maximum(rad, 1e-10)).astype(np.float32)
        y = (2 * area / np.maximum(a*c, 1e-7)).clip(0, 1).astype(np.float32)
        return sides, y[:, None]
    xtr, ytr = sample(n_train, seed)
    xte, yte = sample(n_test, seed + 5000)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": 1}