import numpy as np

META = {
    'name': 'masked_blur_inverse',
    'domain': 'pde_inverse',
    'description': 'Small linear inverse imaging task with spatially nonuniform detector coverage, analogous to a discretized boundary-value inverse problem.'
}

def operator(n=16):
    z = np.arange(n, dtype=np.float64)
    g = np.exp(-((z[:, None] - z[None, :]) ** 2) / (2 * 1.15 ** 2))
    g /= g.sum(axis=1, keepdims=True)
    blur = np.kron(g, g)
    yy, xx = np.mgrid[:n, :n]
    coverage = np.where(xx < int(.58*n), 1.0, .18) + .15
    A = np.sqrt(coverage.reshape(-1))[:, None] * blur
    return A, coverage

def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(seed)
    n = 16
    A, _ = operator(n)
    def sample(count):
        ys = np.zeros((count, n, n), dtype=np.float32)
        yy, xx = np.mgrid[:n, :n]
        for k in range(count):
            cy1, cx1 = rng.integers(3, 13, size=2)
            cy2, cx2 = rng.integers(3, 13, size=2)
            ys[k] = (((yy-cy1)**2 + (xx-cx1)**2) <= 2.0**2).astype(np.float32)
            ys[k] += .9 * (((yy-cy2)**2 + (xx-cx2)**2) <= 1.7**2)
        flat = ys.reshape(count, -1).astype(np.float64)
        clean = flat @ A.T
        noise = .08 * np.std(clean, axis=1, keepdims=True) * rng.normal(size=clean.shape)
        return (clean + noise).astype(np.float32), ys
    xtr, ytr = sample(n_train)
    xte, yte = sample(n_test)
    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
            'task': 'regression', 'metric': 'mse', 'input_shape': (n*n,),
            'out_dim': n*n, 'operator': A.astype(np.float32), 'n': n}