import math
import numpy as np

META = {
    'name': 'symmetric_polynomial_regression',
    'domain': 'symmetric_interactions',
    'description': 'Permutation-invariant degree-four polynomial regression benchmark.'
}

def _compositions(d, m):
    out = []
    def rec(i, left, cur):
        if i == d - 1:
            out.append(tuple(cur + [left]))
        else:
            for k in range(left + 1):
                rec(i + 1, left - k, cur + [k])
    rec(0, m, [])
    return out

def get_dataset(seed, n_train, n_test):
    rng = np.random.RandomState(int(seed))
    d, m = 8, 4
    alphas = _compositions(d, m)
    mult = np.array([math.factorial(m) / np.prod([math.factorial(a) for a in al]) for al in alphas])
    coeff = rng.normal(0, 0.20, len(alphas)) / np.sqrt(mult)
    def make(n):
        x = rng.normal(size=(n, d)).astype(np.float32)
        z = np.ones((n, len(alphas)), dtype=np.float64)
        for k, alpha in enumerate(alphas):
            for j, p in enumerate(alpha):
                if p:
                    z[:, k] *= x[:, j].astype(np.float64) ** p
        y = z.dot(mult * coeff) + rng.normal(0, 0.05, n)
        return x, y.astype(np.float32)[:, None]
    xtr, ytr = make(n_train)
    xte, yte = make(n_test)
    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
            'input_shape': (d,), 'out_dim': 1, 'task': 'regression', 'metric': 'mse'}