import numpy as np

META = {
    'name': 'planar_ising_autoregressive',
    'domain': 'sequence-level sampling',
    'description': 'Planar 4x4 Ising raster-prefix conditional prediction with exact enumerated conditional probabilities.'
}
L, N, BETA = 4, 16, 0.65
EDGES = ([(r*L+c, r*L+c+1) for r in range(L) for c in range(L-1)] +
         [(r*L+c, (r+1)*L+c) for r in range(L-1) for c in range(L)])
_CACHE = None

def _tables():
    global _CACHE
    if _CACHE is not None:
        return _CACHE
    states = np.array(np.meshgrid(*([[-1, 1]] * N))).T.reshape(-1, N).astype(np.int8)
    logw = np.zeros(len(states), dtype=np.float64)
    for u, v in EDGES:
        logw += BETA * states[:, u] * states[:, v]
    logw -= logw.max()
    probs = np.exp(logw); probs /= probs.sum()
    table = {}
    for i in range(N):
        groups = {}
        for j, s in enumerate(states):
            key = tuple(int(v) for v in s[:i])
            z = groups.setdefault(key, [0.0, 0.0])
            z[0 if s[i] < 0 else 1] += probs[j]
        for key, z in groups.items():
            table[(i, key)] = z[1] / (z[0] + z[1])
    _CACHE = (states, probs, table)
    return _CACHE

def get_dataset(seed, n_train, n_test):
    states, probs, table = _tables()
    rng = np.random.RandomState(seed)
    def make(n):
        ix = rng.choice(len(states), size=n, p=probs)
        pos = rng.randint(0, N, size=n)
        x = np.zeros((n, 2*N), dtype=np.float32)
        y = np.zeros((n, 1), dtype=np.float32)
        q = np.zeros((n, 1), dtype=np.float32)
        for k, j in enumerate(ix):
            s = states[j]; i = int(pos[k])
            x[k, :i] = s[:i]; x[k, N+i] = 1.0
            q[k, 0] = table[(i, tuple(int(v) for v in s[:i]))]
            y[k, 0] = float(rng.rand() < q[k, 0])
        return x, y, q
    xtr, ytr, qtr = make(n_train); xte, yte, qte = make(n_test)
    return {'xtr':xtr, 'ytr':ytr, 'xte':xte, 'yte':yte,
            'qtr':qtr, 'qte':qte, 'task':'regression', 'metric':'mse', 'out_dim':1}