"""Volume from the stored atoms in one contraction. The height atoms are orthonormal on a part's layer centres
(Q, _parts.orth_basis), so the sum over the layers of a ring's area pi sum_k k (ax_k by_k - bx_k ay_k) is the same
expression in the atom vectors: V = dz * prod(scale) * pi * sum_parts sign_p |sum_k k (<ax_k, by_k> - <bx_k, ay_k>)|.
A part's rings all run one way, but not always the way its sign says (the last parts of mug ace, shoe, jogger run
reversed): hence sign_p |.| per part. Equal to the layer-centre sum of volume.py (sub=1)."""
import numpy as np

from fcp._slice import centres


def volume_atoms(d):
    K, M = d.K, d.M
    A = np.zeros(len(d.parts) * 2 * K * (M + 1))
    i, v = d.atoms
    A[i] = v
    A = A.reshape(len(d.parts), 2, K, M + 1)  # part, x / y, Fourier term (0, cos 1, sin 1, ...), height atom
    k = np.arange(1, (K - 1) // 2 + 1)
    s = np.einsum("pkm,pkm->pk", A[:, 0, 1::2], A[:, 1, 2::2]) - np.einsum("pkm,pkm->pk", A[:, 0, 2::2], A[:, 1, 1::2])
    sign = np.array([p[0] for p in d.parts], float)
    dz = 2 * float(centres(d.layers)[-1]) / (d.layers - 1)
    return float(np.pi * (sign * np.abs((k * s).sum(1))).sum() * dz * np.prod(np.broadcast_to(np.asarray(d.pose.scale, float), (3,))))


try:
    import numba

    @numba.njit(cache=True)
    def _pairs(i, v, A, sign, P, K, M1):
        """sum_k k (<ax_k, by_k> - <bx_k, ay_k>): the atoms scattered into A (a zeroed scratch, zeroed again on exit),
        then nested loops over part, k and height atom (no integer division: 3 per atom cost 100 us)."""
        for j in range(len(i)):
            A[i[j]] = v[j]
        s = 0.0
        for p in range(P):
            x = p * 2 * K * M1
            y = x + K * M1
            sp = 0.0
            for k in range(1, (K - 1) // 2 + 1):
                c, sn = (2 * k - 1) * M1, 2 * k * M1
                t = 0.0
                for m in range(M1):
                    t += A[x + c + m] * A[y + sn + m] - A[x + sn + m] * A[y + c + m]
                sp += k * t
            s += sign[p] * abs(sp)
        for j in range(len(i)):
            A[i[j]] = 0.0
        return s
except ImportError:
    _pairs = None

_SCRATCH = {}


def volume_atoms_fast(d):
    """volume_atoms in one numba kernel (a scratch array reused per size)."""
    i, v = d.atoms
    P, K, M1 = len(d.parts), d.K, d.M + 1
    A = _SCRATCH.get(P * 2 * K * M1)
    if A is None:
        A = _SCRATCH[P * 2 * K * M1] = np.zeros(P * 2 * K * M1)
    s = _pairs(i, v.astype(np.float64), A, np.array([p[0] for p in d.parts], np.float64), P, K, M1)
    dz = 2 * float(centres(d.layers)[-1]) / (d.layers - 1)
    return float(np.pi * s * dz * np.prod(np.broadcast_to(np.asarray(d.pose.scale, float), (3,))))