# Path: pyem/geom/geom_numba.py
# @numba.jit(nopython=True, cache=False, nogil=True)
# def cross3_sca(u, v):
#     w = np.zeros_like(u)
#     w[0] = u[1] * v[2] - u[2] * v[1]
#     w[1] = u[2] * v[0] - u[0] * v[2]
#     w[2] = u[0] * v[1] - u[1] * v[0]
#     return w


import numba
import numpy as np
from .geom_numba import cross3_sca
    v *= 2
    # msk = v > pi_half
    # v[msk] = np.pi - v[msk]
    return v


@numba.jit(cache=False, nopython=True, parallel=True)
def cdistq(q1, q2, d):
    # if q2 is None:
    #     q2 = q1
    # if out is None:
    # d = np.zeros((len(q1), len(q2)), dtype=q1.dtype)
    # else:
    #     d = out
    #     assert len(d) == len(q1)
    #     assert len(d[0]) == len(q2)
    pi_half = np.pi / 2
    for i in numba.prange(d.shape[0]):
        for j in range(d.shape[1]):
            v = np.abs(np.sum(q1[i] * q2[j]))
            if v > 1.0:
                v = 1.0
            v = np.arccos(v)
            v *= 2
            if v > pi_half:
                v = np.pi - v
            # v += 1e-6
            v **= 2
            v *= -0.5
            d[i, j] = v

# Path: pyem/geom/geom_numba.py
# @numba.jit(nopython=True, cache=False, nogil=True)
# def cross3_sca(u, v):
#     w = np.zeros_like(u)
#     w[0] = u[1] * v[2] - u[2] * v[1]
#     w[1] = u[2] * v[0] - u[0] * v[2]
#     w[2] = u[0] * v[1] - u[1] * v[0]
#     return w


# Path: pyem/geom/quat_numba.py
import numba
import numpy as np
from .geom_numba import cross3_sca

    v *= 2
    # msk = v > pi_half
    # v[msk] = np.pi - v[msk]
    return v


@numba.jit(cache=False, nopython=True, parallel=True)
def cdistq(q1, q2, d):
    # if q2 is None:
    #     q2 = q1
    # if out is None:
    # d = np.zeros((len(q1), len(q2)), dtype=q1.dtype)
    # else:
    #     d = out
    #     assert len(d) == len(q1)
    #     assert len(d[0]) == len(q2)
    pi_half = np.pi / 2
    for i in numba.prange(d.shape[0]):
        for j in range(d.shape[1]):
            v = np.abs(np.sum(q1[i] * q2[j]))
            if v > 1.0:
                v = 1.0
            v = np.arccos(v)
            v *= 2
            if v > pi_half:
                v = np.pi - v
            # v += 1e-6
            v **= 2
            v *= -0.5
            d[i, j] = v

    return d