import numpy as np

META = {
    "name": "harmonic_cycle_circulation",
    "domain": "pde",
    "description": "Regression on a discretized annulus-like one-cycle: a local source encodes randomized global circulation and the target is the corresponding harmonic edge 1-form."
}


def harmonic_basis(n=24):
    # Vertices/edges of a periodic annulus skeleton. B1 is vertex-edge incidence;
    # there are no filled 2-cells in this minimal annulus skeleton.
    B1 = np.zeros((n, n), dtype=np.float64)
    for e in range(n):
        B1[e, e] = -1.0
        B1[(e + 1) % n, e] = 1.0
    masses = 0.5 + np.linspace(0.2, 1.4, n)
    M = np.diag(masses)
    h = np.linalg.solve(M, np.ones(n))
    h /= np.sqrt(h @ M @ h)
    return B1, M, h[:, None]


def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.RandomState(seed)
    _, _, H = harmonic_basis()
    atr = rng.uniform(-2.0, 2.0, n_train).astype(np.float32)
    ate = np.random.RandomState(seed + 5000).uniform(-2.0, 2.0, n_test).astype(np.float32)
    xtr = np.zeros((n_train, len(H), 1), dtype=np.float32)
    xte = np.zeros((n_test, len(H), 1), dtype=np.float32)
    # A source is local (one edge); four graph layers cannot directly reach the
    # opposite side, while global pooling can encode the circulation coefficient.
    xtr[:, 0, 0] = atr
    xte[:, 0, 0] = ate
    ytr = atr[:, None] * H[:, 0][None, :].astype(np.float32)
    yte = ate[:, None] * H[:, 0][None, :].astype(np.float32)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": len(H)}