import numpy as np
import torch

META = {
    "name": "incidence_actions",
    "domain": "structured_action_classification",
    "description": "Variable-size atomic entities and legal composite actions represented by an instance-specific binary incidence matrix; labels are noisy utility-maximizing legal actions."
}
D = 4
MAX_ATOMS = 24
MAX_ACTIONS = 48

def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(int(seed))
    true_w = np.array([1.2, -0.9, 0.65, 0.35], dtype=np.float32)
    def make(n):
        X = np.zeros((n, MAX_ATOMS * D + MAX_ACTIONS * MAX_ATOMS), dtype=np.float32)
        y = np.zeros(n, dtype=np.int64)
        for b in range(n):
            na = int(rng.integers(5, MAX_ATOMS + 1))
            pa = int(rng.integers(5, MAX_ACTIONS + 1))
            atoms = rng.normal(size=(na, D)).astype(np.float32)
            A = np.zeros((pa, na), dtype=np.float32)
            for p in range(pa):
                k = int(rng.integers(1, min(6, na) + 1))
                A[p, rng.choice(na, size=k, replace=False)] = 1.0
            utilities = A @ (atoms @ true_w) + 0.08 * rng.normal(size=pa)
            y[b] = int(np.argmax(utilities))
            X[b, :na * D] = atoms.reshape(-1)
            packed = np.pad(A, ((0, 0), (0, MAX_ATOMS-na))).reshape(-1)
            X[b, MAX_ATOMS * D:MAX_ATOMS * D + pa * MAX_ATOMS] = packed
        return X, y
    xtr, ytr = make(n_train)
    xte, yte = make(n_test)
    return {"xtr": torch.from_numpy(xtr), "ytr": torch.from_numpy(ytr),
            "xte": torch.from_numpy(xte), "yte": torch.from_numpy(yte),
            "input_shape": (xtr.shape[1],), "out_dim": MAX_ACTIONS,
            "task": "classification", "metric": "cross_entropy"}