import numpy as np

META = {
    "name": "masked_candidate_low_rank",
    "domain": "retrieval",
    "description": "Masked high-dimensional candidate vectors generated from a low-rank latent action-feature subspace; regression predicts candidate utility."
}


def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.RandomState(seed)
    d, rank = 128, 8
    q, _ = np.linalg.qr(rng.normal(size=(d, rank)))
    w = rng.normal(size=rank)

    def make(n, offset):
        rr = np.random.RandomState(seed + offset)
        z = rr.normal(size=(n, rank)).astype(np.float32)
        x = (z @ q.T).astype(np.float32)
        # A fixed latent utility with modest observation-independent noise.
        y = (z @ w + 0.15 * np.sin(z[:, 0]) + 0.08 * rr.normal(size=n)).astype(np.float32)
        mask = (rr.rand(n, d) < 0.5).astype(np.float32)
        return x, y, mask

    xtr, ytr, mtr = make(n_train, 17)
    xte, yte, mte = make(n_test, 50017)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "mtr": mtr, "mte": mte, "task": "regression", "metric": "mse",
            "out_dim": 1, "rank": rank, "p": 0.5}