import numpy as np
from scipy.spatial.transform import Rotation as R, Slerp

def aa_to_quat(aa: np.ndarray) -> np.ndarray:
    flat = aa.reshape(-1, 3)
    q = R.from_rotvec(flat).as_quat()           # xyzw
    return q.reshape(*aa.shape[:-1], 4)

def fix_hemisphere(quats: np.ndarray) -> np.ndarray:
    out = quats.copy()
    for i in range(1, len(out)):
        if np.dot(out[i], out[i - 1]) < 0:
            out[i] = -out[i]
    return out

def slerp_series(quats: np.ndarray, src_t: np.ndarray, dst_t: np.ndarray) -> np.ndarray:
    if len(quats) == 1:
        return np.repeat(quats, len(dst_t), axis=0)
    slerp = Slerp(src_t, R.from_quat(quats))
    clamped = np.clip(dst_t, src_t[0], src_t[-1])
    return slerp(clamped).as_quat()