import numpy as np

META = {
    "name": "orientation_phase_quotient",
    "domain": "molecular_orientation_symmetry",
    "description": "Cubic apolar molecular orientations with isotropic/nematic phase labels and nuisance global rotations/sign flips.",
}


def _p2(x):
    return 0.5 * (3.0 * x * x - 1.0)


def _rot(rng):
    q, r = np.linalg.qr(rng.normal(size=(3, 3)))
    q = q @ np.diag(np.where(np.diag(r) >= 0, 1.0, -1.0))
    if np.linalg.det(q) < 0:
        q[:, 0] *= -1
    return q


def _corr_features(u, L=4, shells=2):
    a = u.reshape(L, L, L, 3)
    channels = []
    for radius in range(1, shells + 1):
        offsets = [(dx, dy, dz)
                   for dx in range(-radius, radius + 1)
                   for dy in range(-radius, radius + 1)
                   for dz in range(-radius, radius + 1)
                   if (dx, dy, dz) != (0, 0, 0)
                   and dx * dx + dy * dy + dz * dz == radius * radius]
        c = np.zeros((L, L, L), dtype=np.float64)
        for dx, dy, dz in offsets:
            b = np.roll(a, (dx, dy, dz), axis=(0, 1, 2))
            c += _p2(np.sum(a * b, axis=-1))
        channels.append(c / len(offsets))
    return np.stack(channels, axis=0).reshape(-1).astype(np.float32)


def _sample(rng, phase, L=4):
    n = L ** 3
    if phase == 0:
        u = rng.normal(size=(n, 3))
    else:
        d = rng.normal(size=3)
        d /= np.linalg.norm(d)
        u = d + rng.normal(scale=0.28, size=(n, 3))
    u /= np.linalg.norm(u, axis=1, keepdims=True)
    u *= rng.choice([-1.0, 1.0], size=(n, 1))
    # Global frame is nuisance; independent signs are the apolar gauge.
    return u @ _rot(rng).T


def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(int(seed))
    total = int(n_train) + int(n_test)
    x, y = [], []
    for k in range(total):
        phase = k % 2
        x.append(_sample(rng, phase))
        y.append(phase)
    x = np.asarray(x, dtype=np.float32)
    y = np.asarray(y, dtype=np.int64)
    perm = rng.permutation(total)
    x, y = x[perm], y[perm]
    return {
        "xtr": x[:n_train].reshape(n_train, -1),
        "ytr": y[:n_train],
        "xte": x[n_train:].reshape(n_test, -1),
        "yte": y[n_train:],
        "task": "classification",
        "metric": "err",
        "input_shape": (4 * 4 * 4 * 3,),
        "out_dim": 2,
    }


def invariant_dataset(ds):
    def conv(x):
        return np.asarray([_corr_features(v.reshape(4 ** 3, 3)) for v in x], dtype=np.float32)
    out = dict(ds)
    out["xtr"] = conv(ds["xtr"])
    out["xte"] = conv(ds["xte"])
    out["input_shape"] = (2 * 4 ** 3,)
    return out


def math_check(seed=71):
    rng = np.random.default_rng(seed)
    u = _sample(rng, 1)
    v = u @ _rot(rng).T
    v *= rng.choice([-1.0, 1.0], size=(len(v), 1))
    return float(np.max(np.abs(_corr_features(u) - _corr_features(v))))