import numpy as np

META = {
    "name": "semantic_pushforward_classification",
    "domain": "uncertainty_calibration",
    "description": "Synthetic finite-state classification with multiple verbal continuations per state and controlled response-logit distortion; evaluates calibrated semantic pushforward."
}


def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(int(seed))
    k, d = 3, 8
    W = rng.normal(0, 0.8, (k, d))

    def sample(n):
        x = rng.normal(size=(n, d)).astype(np.float32)
        z = x @ W.T + 0.45 * np.sin(x[:, :3] @ W[:, :3].T)
        q = np.exp(z - z.max(1, keepdims=True))
        q /= q.sum(1, keepdims=True)
        y = np.array([rng.choice(k, p=p) for p in q], dtype=np.int64)
        return x, y

    xtr, ytr = sample(int(n_train))
    xte, yte = sample(int(n_test))
    return {
        "xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
        "task": "classification", "metric": "error",
        "out_dim": k, "input_shape": (d,)
    }