# Path: pyem/geom/quat_numba.py
# @numba.jit(cache=False, nopython=True, nogil=True)
# def distq(q1, q2):
#     pi_half = np.pi / 2
#     v = np.abs(np.sum(q1 * q2))
#     v = np.clip(0, 1.0, v)
#     v = np.arccos(v)
#     v *= 2
#     # msk = v > pi_half
#     # v[msk] = np.pi - v[msk]
#     return v
# 
# Path: pyem/geom/quat_numba.py
# @numba.jit(cache=True, nopython=True, nogil=True)
# def qslerp(q1, q2, t, longest=False):
#     cos_half_theta = np.dot(q1, q2)
#     if cos_half_theta >= 1.0:
#         return q1.copy()
#     if longest:
#         if cos_half_theta > 0:
#             cos_half_theta = -cos_half_theta
#             q1 = -q1
#     elif cos_half_theta < 0:
#         cos_half_theta = -cos_half_theta
#         q1 = -q1
#     half_theta = np.arccos(cos_half_theta)
#     sin_half_theta = np.sqrt(1 - cos_half_theta * cos_half_theta)
#     if np.abs(sin_half_theta) < 1E-12:
#         return (q1 + q2) / 2
#     a = np.sin((1 - t) * half_theta)
#     b = np.sin(t * half_theta)
#     return (q1 * a + q2 * b) / sin_half_theta
# 
# Path: pyem/geom/quat_numba.py
# @numba.guvectorize(["void(float64[:], float64[:], float64[:])"],
#                    "(m),(m)->(m)", nopython=True, cache=False)
# def qtimes(q1, q2, q3):
#     _qtimes(q1, q2, q3)
# 
# Path: pyem/geom/quat.py
# def meanq(q, w=None):
#     if w is None:
#         return np.linalg.eigh(np.einsum('ij,ik->...jk', q, q))[1][:, -1]
#     else:
#         return np.linalg.eigh(np.einsum('ij,ik,i->...jk', q, q, w))[1][:, -1]


import numpy as np
from .quat_numba import distq
from .quat_numba import qslerp
from .quat_numba import qtimes
from .quat import meanq
    qexp = []
    for i in range(0, len(keyq) - 1):
        ninterp = np.round(np.rad2deg(2 * distq(keyq[i], keyq[i + 1])) * steps_per_deg)
        for t in np.arange(0, 1, 1 / ninterp):
            qexp.append(qslerp(keyq[i], keyq[i + 1], t))
    qexp.append(keyq[-1])
    return np.array(qexp)


def findkeyq(qarr, kpcs, nkey=10, pc_cyl_ptile=25, pc_ptile=99, pc=0):
    otherpcs = list(set(np.arange(kpcs.shape[1])) - {pc})
    pcr = np.sqrt(np.sum(kpcs[:, otherpcs] ** 2, axis=1))
    mask = (pcr < np.percentile(pcr, pc_cyl_ptile)) & (
                np.abs(kpcs[:, pc]) < np.percentile(np.abs(kpcs[:, pc]), pc_ptile))
    mn, mx = np.min(kpcs[mask, pc]), np.max(kpcs[mask, pc])
    bins = np.linspace(mn, mx, nkey)
    idx = np.digitize(kpcs[mask, pc], bins) - 1
    uidx, uidxcnt = np.unique(idx, return_counts=True)
    keyq = []
    for i in uidx[uidxcnt >= 10]:
        kq = meanq(qarr[mask, :][idx == i, :])
        keyq.append(kq)
    return np.array(keyq)


def dualquat(q, t):
    assert q.shape[0] == t.shape[0]
    dq = np.zeros(q.shape, dtype=np.complex128)
    dq.real = q
    dq.imag[:, 1:] = t

# Path: pyem/geom/quat_numba.py
# @numba.jit(cache=False, nopython=True, nogil=True)
# def distq(q1, q2):
#     pi_half = np.pi / 2
#     v = np.abs(np.sum(q1 * q2))
#     v = np.clip(0, 1.0, v)
#     v = np.arccos(v)
#     v *= 2
#     # msk = v > pi_half
#     # v[msk] = np.pi - v[msk]
#     return v
# 
# Path: pyem/geom/quat_numba.py
# @numba.jit(cache=True, nopython=True, nogil=True)
# def qslerp(q1, q2, t, longest=False):
#     cos_half_theta = np.dot(q1, q2)
#     if cos_half_theta >= 1.0:
#         return q1.copy()
#     if longest:
#         if cos_half_theta > 0:
#             cos_half_theta = -cos_half_theta
#             q1 = -q1
#     elif cos_half_theta < 0:
#         cos_half_theta = -cos_half_theta
#         q1 = -q1
#     half_theta = np.arccos(cos_half_theta)
#     sin_half_theta = np.sqrt(1 - cos_half_theta * cos_half_theta)
#     if np.abs(sin_half_theta) < 1E-12:
#         return (q1 + q2) / 2
#     a = np.sin((1 - t) * half_theta)
#     b = np.sin(t * half_theta)
#     return (q1 * a + q2 * b) / sin_half_theta
# 
# Path: pyem/geom/quat_numba.py
# @numba.guvectorize(["void(float64[:], float64[:], float64[:])"],
#                    "(m),(m)->(m)", nopython=True, cache=False)
# def qtimes(q1, q2, q3):
#     _qtimes(q1, q2, q3)
# 
# Path: pyem/geom/quat.py
# def meanq(q, w=None):
#     if w is None:
#         return np.linalg.eigh(np.einsum('ij,ik->...jk', q, q))[1][:, -1]
#     else:
#         return np.linalg.eigh(np.einsum('ij,ik,i->...jk', q, q, w))[1][:, -1]


# Path: pyem/geom/geom.py
import numpy as np
from .quat_numba import distq
from .quat_numba import qslerp
from .quat_numba import qtimes
from .quat import meanq

    qexp = []
    for i in range(0, len(keyq) - 1):
        ninterp = np.round(np.rad2deg(2 * distq(keyq[i], keyq[i + 1])) * steps_per_deg)
        for t in np.arange(0, 1, 1 / ninterp):
            qexp.append(qslerp(keyq[i], keyq[i + 1], t))
    qexp.append(keyq[-1])
    return np.array(qexp)


def findkeyq(qarr, kpcs, nkey=10, pc_cyl_ptile=25, pc_ptile=99, pc=0):
    otherpcs = list(set(np.arange(kpcs.shape[1])) - {pc})
    pcr = np.sqrt(np.sum(kpcs[:, otherpcs] ** 2, axis=1))
    mask = (pcr < np.percentile(pcr, pc_cyl_ptile)) & (
                np.abs(kpcs[:, pc]) < np.percentile(np.abs(kpcs[:, pc]), pc_ptile))
    mn, mx = np.min(kpcs[mask, pc]), np.max(kpcs[mask, pc])
    bins = np.linspace(mn, mx, nkey)
    idx = np.digitize(kpcs[mask, pc], bins) - 1
    uidx, uidxcnt = np.unique(idx, return_counts=True)
    keyq = []
    for i in uidx[uidxcnt >= 10]:
        kq = meanq(qarr[mask, :][idx == i, :])
        keyq.append(kq)
    return np.array(keyq)


def dualquat(q, t):
    assert q.shape[0] == t.shape[0]
    dq = np.zeros(q.shape, dtype=np.complex128)
    dq.real = q
    dq.imag[:, 1:] = t

    dq.imag = qtimes(dq.imag, dq.real) * 0.5