import numpy as np

META = {
    "name": "complex_array_channel",
    "domain": "array-valued complex low-rank tensors",
    "description": "Regression from noisy complex array channels to clean channels represented by continuous rank-one steering atoms."
}

def steering(n, u):
    x = np.arange(n, dtype=np.float32) - (n - 1) / 2.0
    return np.exp(1j * np.pi * x * u) / np.sqrt(n)

def make_channel(rng, n=8, k=2):
    h = np.zeros((n, n), dtype=np.complex64)
    for _ in range(k):
        ur, ut = rng.uniform(-0.8, 0.8, 2)
        g = (rng.normal() + 1j * rng.normal()) / np.sqrt(2 * k)
        h += g * np.outer(steering(n, ur), steering(n, ut).conj())
    return h

def pack(h):
    return np.concatenate([h.real.reshape(-1), h.imag.reshape(-1)]).astype(np.float32)

def get_dataset(seed, n_train, n_test):
    rng = np.random.RandomState(seed)
    def sample(num):
        xs, ys = [], []
        for _ in range(num):
            clean = make_channel(rng)
            noise = (rng.normal(size=clean.shape) + 1j * rng.normal(size=clean.shape)).astype(np.complex64) * 0.035
            xs.append(pack(clean + noise))
            ys.append(pack(clean))
        return np.asarray(xs, dtype=np.float32), np.asarray(ys, dtype=np.float32)
    xtr, ytr = sample(n_train)
    xte, yte = sample(n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse"}