import numpy as np

META = {
    'name': 'tetra_elasticity_complex',
    'domain': 'pde',
    'description': 'Simplicial tetrahedral edge-field regression with vertex-edge-face-cell incidence operators.'
}

TETS = ((0, 1, 2, 3), (0, 2, 1, 4))


def operators():
    edges = sorted({tuple(sorted((t[i], t[j]))) for t in TETS
                    for i in range(4) for j in range(i + 1, 4)})
    faces = sorted({tuple(sorted(t[j] for j in range(4) if j != omit))
                    for t in TETS for omit in range(4)})
    ei, fi = {e: i for i, e in enumerate(edges)}, {f: i for i, f in enumerate(faces)}
    D0 = np.zeros((len(edges), 5), dtype=np.float32)
    for r, (a, b) in enumerate(edges):
        D0[r, a], D0[r, b] = -1, 1
    D1 = np.zeros((len(faces), len(edges)), dtype=np.float32)
    for r, (a, b, c) in enumerate(faces):
        for e, s in (((b, c), 1), ((a, c), -1), ((a, b), 1)):
            D1[r, ei[e]] = s
    D2 = np.zeros((len(TETS), len(faces)), dtype=np.float32)
    for r, t in enumerate(TETS):
        for omit in range(4):
            sub = tuple(t[j] for j in range(4) if j != omit)
            f = tuple(sorted(sub))
            inv = sum(sub[i] > sub[j] for i in range(3) for j in range(i + 1, 3))
            D2[r, fi[f]] = (-1) ** omit * (-1 if inv % 2 else 1)
    return D0, D1, D2, edges, faces

D0, D1, D2, EDGES, FACES = operators()


def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(int(seed))
    # Each sample is an edge 1-cochain: compatible gradient + optional face defect.
    # Target is the observed incompatibility magnitude, an ordinary regression MSE.
    def make(n):
        u = rng.normal(size=(n, 5)).astype(np.float32)
        x = u @ D0.T
        defect = rng.normal(size=(n, len(EDGES))).astype(np.float32)
        active = (rng.random(n) < 0.65).astype(np.float32)[:, None]
        x = x + active * 0.30 * defect
        y = np.sqrt(np.mean((x @ D1.T) ** 2, axis=1, keepdims=True)).astype(np.float32)
        return x.astype(np.float32), y
    xtr, ytr = make(n_train)
    xte, yte = make(n_test)
    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
            'task': 'regression', 'metric': 'mse', 'input_shape': (len(EDGES),),
            'out_dim': 1, 'D0': D0, 'D1': D1, 'D2': D2}