import numpy as np

META = {
    'name': 'bidirectional_conditional_joint',
    'domain': 'conditional_compatibility',
    'description': 'Strictly positive categorical joint samples for two neural conditional directions.'
}

def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.default_rng(int(seed))
    raw = np.array([
        [1.0, 2.0, .7, 1.4, .8, 1.7],
        [2.1, .8, 1.8, .5, 1.3, .9],
        [.6, 1.7, 2.4, 1.1, .9, 1.5],
        [1.3, .9, 1.5, 2.2, 1.1, .7],
        [1.8, 1.2, .5, 1.6, 2.0, 1.0],
        [.7, 1.4, 1.1, .8, 1.6, 2.3]], dtype=np.float64)
    joint = raw / raw.sum()
    z = rng.choice(36, size=n_train + n_test, p=joint.reshape(-1))
    pairs = np.stack([z // 6, z % 6], axis=1).astype(np.int64)
    return {
        'xtr': pairs[:n_train], 'ytr': pairs[:n_train, 0],
        'xte': pairs[n_train:], 'yte': pairs[n_train:, 0],
        'task': 'classification', 'metric': 'bidirectional_nll',
        'input_shape': (2,), 'out_dim': 6, 'nx': 6, 'ny': 6,
        'joint': joint
    }