import numpy as np

META = {
    "name": "persistence_diagrams",
    "domain": "topological_representation",
    "description": "Classification of noisy variable-cardinality persistence diagrams; class is encoded by interval location and persistence."
}


def _make(seed, n):
    rng = np.random.default_rng(seed)
    ys = rng.integers(0, 2, size=n, dtype=np.int64)
    out = np.zeros((n, 8, 2), dtype=np.float32)
    for i, y in enumerate(ys):
        m = int(rng.integers(2, 9))
        center = 0.29 if y == 0 else 0.69
        births = np.clip(rng.normal(center, 0.105, m), 0.03, 0.86)
        persistence = np.clip(rng.normal(0.20, 0.065, m), 0.035, 0.38)
        deaths = np.minimum(births + persistence, 0.97)
        deaths = np.maximum(deaths, births + 0.015)
        out[i, :m, 0] = births
        out[i, :m, 1] = deaths
    return out, ys


def get_dataset(seed, n_train, n_test):
    xtr, ytr = _make(int(seed) * 2 + 11, n_train)
    xte, yte = _make(int(seed) * 2 + 10011, n_test)
    return {
        "xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
        "task": "classification", "metric": "err",
        "input_shape": xtr.shape[1:], "out_dim": 2
    }