"""
Provides a rotation class and utils.

Implements different algorithms and optimisations.

Currently, QERRotator seems best, closely followed by ERRotator.

QuatRotator seems to come in next, though it is a bit clunky code due
to limited support for `.imag` on numpy ndarrays and functions to
create quaternions with just imaginary part with a direct helper
function.  Should write own based in fortran and f2py;
we only use small subset of quaternion package.

(C) Alexander Heger 2019, 2021--2023
"""

from itertools import permutations
from multiprocessing import Pool

import numpy as np
from numpy import linalg as LA
from scipy.linalg import expm
try:
    import quaternion as quat
except:
    raise Exception(f'''
    [{__name__}] ERROR: require module "numpy-quaternion".
    [{__name__}]
    [{__name__}] To do this, please esecute, e.g., on the command line
    [{__name__}]
    [{__name__}] pip3 install -U numpy-quaternion
    [{__name__}] ''')
import math
try:
    import numba
except:
    # print(f' [{__name__}] WARNING: module numba not found.')
    class numba:
        @staticmethod
        def jit(*args, **kwargs):
            def no_jit(f):
                # print(f' [{__name__}] WARNING: not optimising "{f.__name__}".')
                return f
            return no_jit

from .const import *
from .util import cross_vec, cross, norm, dot, norm2

class NoRotator(object):
    def __init__(self, *args, **kwargs):
        pass
    def rotate(self, v, *args, **kwargs):
        return v
    def __call__(self, *args, **kwargs):
        return self.rotate(*args, **kwargs)
    def update(self, *args, **kwargs):
        pass


class BaseRotator(object):
    """
    Basic Rotator functionallities.

    Allow np.ndarrays for time, rotation vectors, and vectors.  For
    the latter two, it is assume that the vector 3-dimension is the
    last one of the array.

    The return array will first have time dimensions, then rotator
    dimensions, and last vector dimension, terminated by the final
    3-dimension.

    The `align` parmeter allwos vectorisation of of the first `align`
    dimensions of the vector array with the last `align` non-vector
    dimensions of the omega array.

    The `phase` parameter allows to add a phase offset to the rotation
    matices.  It needs to be broadcastbale to the non-vector dimensions
    of the omega array.  Phase is 2 pi for a full rotation (I think).

    TODO: Use tuples to allow alignment with arbitraty axes, and with
    time axes.
    """

    void = object()

    def __init__(self, omega, time = None, phase = None):
        self.time = None
        self.phase = None

        if isinstance(omega, np.ndarray) and (omega.dtype == quat.quaternion):
            self.initq(omega)
        elif np.shape(omega)[-1] == 4:
            self.initf(omega)
        else:
            self.initomega(omega)
        self.init()
        self.settimeandphase(time, phase)

    def init(self):
        pass

    def initomega(self, omega):
        self.omega = np.asarray(omega, dtype = np.float64)

    def initf(self, f):
        self.initq(f2q(f))

    def initq(self, q):
        self.omega = q2w(q)

    def settime(self, time):
        return self.settimeandphase(time = time, phase = None)

    def setphase(self, phase):
        return self.settimeandphase(time = None, phase = phase)

    def settimeandphase(self, time = None, phase = None):
        if time is not None:
            self.time = np.asarray(time)
        if phase is not None:
            phase = np.asarray(phase)
            self.phase = phase / LA.norm(self.omega, axis = -1)
        if ((time is not None or phase is not None) and
            (time is not None or self.time is not None)):
            self.update()
        return self

    def update(self):
        raise NotImplementedError()

    def rotate(self, v):
        raise NotImplementedError()

    def __call__(self, v, time = None, phase = None, align = 0):
        if self.time is None and time is None:
            time = 1
        if time is not None or phase is not None:
            self.settimeandphase(time, phase)
        return self.rotate(v, align)

    def rotate_matrix(self, m, *args, **kwargs):
        assert m.shape[-2:] == (3,3)
        raise NotImplementedError()

class QuatRotator2(BaseRotator):
    """
    This uses the rotation functions from the quaternion module,
    but this is rather slow.
    """
    neutral = quat.one

    def init(self):
        self.rotation = np.tile(self.neutral, self.omega.shape[:-1])

    def update(self):
        wt = np.multiply.outer(self.time, self.omega)
        if self.phase is not None:
            wt += self.omega * self.phase[..., np.newaxis]
        self.rotation = quat.from_rotation_vector(wt)

    def rotate(self, v, align = 0, phase = None):
        if align == 0:
            return quat.rotate_vectors(self.rotation, v, axis=-1)

        assert self.rotation[0].shape[-align:] == v.shape[:align]
        # need to be iterated manually.
        dr = self.rotation.shape
        dv = v.shape
        nr = len(dr)
        nv = len(dv)

        ia = np.prod(dv[:align], dtype=np.int64)
        sr = [np.prod(dr[:nr-align], dtype=np.int64), ia]
        sv = [ia, np.prod(dv[align:-1], dtype=np.int64), 3]
        so = [sr[0]] + sv

        o = np.empty(so)
        vl = v.reshape(sv)
        rl = self.rotation.reshape(sr)

        # this may be a case for numba?
        for i in range(ia):
            o[:,i] = quat.rotate_vectors(rl[:,i], vl[i], axis=-1)

        do = dr + dv[align:]
        return o.reshape(do)

class QuatRotator(BaseRotator):
    """
    Uses quaternions for rotation.

    Generaotor, g = [rotation verctor] / 2
    Rotator q = exp(g * t)

    Rotated vecor is v' = q * v * q.conjugate()

    where both rotation vector and v components go into the imaginary part.
    """
    neutral = (quat.one,) * 2

    def init(self):
        self.rotation = self.neutral
        f = np.zeros(self.omega.shape[:-1] + (4,))
        f[...,1:] = self.omega[...,:] * 0.5
        self.generator = quat.from_float_array(f)

    def update(self):
        q = np.multiply.outer(self.time, self.generator)
        if self.phase is not None:
            q += self.generator * self.phase
        q = np.exp(q)
        self.rotation = (q, q.conj())

    def rotate(self, v, align = 0, phase = None):
        if align > 0:
            assert self.rotation[0].shape[-align:] == v.shape[:align]

        v = np.asarray(v)
        q = np.zeros(v.shape[:-1] + (4,))
        q[...,1:] = v[...,:]
        q = quat.from_float_array(q)

        dr = self.rotation[0].shape
        dq = q.shape
        nr = len(dr)
        nq = len(dq)
        sr = dr + (1,) * (nq - align)
        sq = (1,) * (nr - align) + dq

        return quat.as_float_array(
            self.rotation[0].reshape(sr)
            * q.reshape(sq)
            * self.rotation[1].reshape(sr)
            )[...,1:]

class QuatCrossRotator(QuatRotator):
    """
    Uses quaternions for rotation (see base class) but executes
    rotation as cross products, supposedly 30% faster.

    There is some replication ...
    """
    def update(self):
        q = np.multiply.outer(self.time, self.generator)
        if self.phase is not None:
            q += self.generator * self.phase
        q = quat.as_float_array(np.exp(q))
        self.rotation = (q[...,0], q[...,1:])

    def rotate(self, v, align = 0, phase = None):
        if align > 0:
            assert self.rotation[0].shape[-align:] == v.shape[:align]

        v = np.asarray(v)

        dq = self.rotation[0].shape
        dv = v.shape[:-1]
        nq = len(dq)
        nv = len(dv)
        sq = dq + (1,) * (nv - align)
        sv = (1,) * (nq - align) + dv + (3,)

        qv = self.rotation[1].reshape(sq + (3,))
        q0 = self.rotation[0].reshape(sq + (1,))
        v = v.reshape(sv)

        # t = np.cross(2 * qv, v, axis=-1)
        # return v + q0 * t + np.cross(qv, t, axis=-1)
        t = cross_vec(2 * qv, v, axis=-1)
        return v + q0 * t + cross_vec(qv, t, axis=-1)


class BaseMatrixRotator(BaseRotator):
    """
    standard matrix multiplication version
    """
    neutral = np.eye(3)

    def init(self):
        self.rotation = np.tile(self.neutral, self.omega.shape[:-1])

    def rotate(self, v, align = 0):
        v = np.asarray(v)
        if align > 0:
            assert self.rotation.shape[-align-2:-2] == v.shape[:align]
        nr = len(self.rotation.shape) - 2
        nv = len(v.shape) - 1
        ir = list(range(2, 2 + nr))
        iv = list(range(2 + nr - align, 2 + nr - align + nv))
        iv[0:align] = ir[len(ir)-align:]
        io = ir + iv[align:] + [0]
        ir += [0, 1]
        iv += [1]
        return np.einsum(self.rotation, ir, v, iv, io, optimize=True)

# numba does not do anything here ...
# @numba.jit(nopython=True, parallel=True)
# def _update_MR(a, b, n):
#     for i in range(n):
#         b[i] = expm(a[i])
# maybe try pool ...
def _update_MR(a, b, n):
    with Pool() as pool:
        b = pool.map(expm, a)


class MatrixRotator(BaseMatrixRotator):
    """
    use generator and matrix exponnent (expm)
    """

    _idxg1 = np.array([2,1,0,2,1,0])
    _idxg2 = np.array([1,2,2,0,0,1])
    _idxo1 = np.array([0,0,1,1,2,2])
    _s1 = np.array([1,-1,1,-1,1,-1])
    _d = (-1,3,3)

    def init(self):
        super().init()
        omega = self.omega
        g = np.zeros(omega.shape[:-1] + (3,3))
        # g[...,2,1] = +omega[...,0]
        # g[...,1,2] = -omega[...,0]
        # g[...,0,2] = +omega[...,1]
        # g[...,2,0] = -omega[...,1]
        # g[...,1,0] = +omega[...,2]
        # g[...,0,1] = -omega[...,2]
        g[..., self._idxg1, self._idxg2] = self._s1 * omega[..., self._idxo1]
        self.generator = g

    def update(self):
        # if len(np.shape(self.time)) == 0 and len(np.shape(self.omega)) == 1:
        #     a = self.generator * self.time
        #     self.rotation = expm(a)
        #     return
        a = np.multiply.outer(self.time, self.generator)
        if self.phase is not None:
            a += self.generator * self.phase[..., np.newaxis, np.newaxis]
        b = np.empty_like(a).reshape(self._d)
        _update_MR(a.reshape(self._d), b, np.prod(a.shape[:-2], dtype = np.int64))
        self.rotation =  b.reshape(a.shape)

        # # alternative naive code is about 1% slower for big arrays
        # self.rotation =  np.asarray(
        #     [expm(x) for x in np.reshape(a, (-1,3,3))]).reshape(a.shape)

@numba.jit(nopython=True)
def _update_ER_scalar(omega, time):
    w = time * omega
    x = LA.norm(w) * 0.5
    a = np.cos(x)
    if x < 1.e-99:
        x = 1.e-99
    y = 0.5 / x
    b, c, d = w * (np.sin(x) * y)
    aa, bb, cc, dd = a * a, b * b, c * c, d * d
    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
    return np.array(
        [[aa + bb - cc - dd, 2 * (bc - ad), 2 * (bd + ac)],
         [2 * (bc + ad), aa + cc - bb - dd, 2 * (cd - ab)],
         [2 * (bd - ac), 2 * (cd + ab), aa + dd - bb - cc]])

# numba makes things slower for large arrays
# @numba.jit(parallel=True)
def _update_ER(omega, time, phase):
    w = np.multiply.outer(time, omega)
    if phase is not None:
        w += omega * phase[..., np.newaxis]
    w = np.moveaxis(w, -1, 0)
    x = LA.norm(w, axis = 0)
    x = np.maximum(x, 1.e-99)
    y = 1 / x
    b, c, d = w * (np.sin(x * 0.5) * y)
    a = np.cos(x * 0.5)
    aa, bb, cc, dd = a * a, b * b, c * c, d * d
    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
    return np.array(
        [[aa + bb - cc - dd, 2 * (bc - ad), 2 * (bd + ac)],
         [2 * (bc + ad), aa + cc - bb - dd, 2 * (cd - ab)],
         [2 * (bd - ac), 2 * (cd + ab), aa + dd - bb - cc]])

class ERRotator(BaseMatrixRotator):
    """Direct computation using Euler-Rodrigues formula"""

    def update(self):
        if len(np.shape(self.time)) == 0 and len(np.shape(self.omega)) == 1:
            if self.phase is None:
                self.rotation = _update_ER_scalar(self.omega, self.time)
                return
            self.rotation = _update_ER_scalar(self.omega, self.time + self.phase)
            return
        m = _update_ER(self.omega, self.time, self.phase)
        self.rotation = np.moveaxis(m, (0, 1), (-2,-1))

# numba makes things slower for large arrays
# @numba.jit(parallel=True)
def _update_ERQ(q):
    a, b, c, d = np.moveaxis(quat.as_float_array(q), -1, 0)
    aa, bb, cc, dd = a * a, b * b, c * c, d * d
    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
    return np.array(
        [[aa + bb - cc - dd, 2 * (bc - ad), 2 * (bd + ac)],
         [2 * (bc + ad), aa + cc - bb - dd, 2 * (cd - ab)],
         [2 * (bd - ac), 2 * (cd + ab), aa + dd - bb - cc]])

class QERRotator(BaseMatrixRotator):
    """
    Use quaternions to compute rotation, then apply rotation using matrix.
    """
    neutral = np.eye(3)

    def initq(self, q):
        super().initq(q)
        self.time = 1
        m = _update_ERQ(q)
        self.rotation = np.moveaxis(m, (0, 1), (-2,-1))

    def init(self):
        q = np.zeros(self.omega.shape[:-1] + (4,))
        q[...,1:] = self.omega[...,:] * 0.5
        self.generator = quat.from_float_array(q)

    def update(self):
        q = np.multiply.outer(self.time, self.generator)
        if self.phase is not None:
            q += self.generator * self.phase
        m = _update_ERQ(np.exp(q))
        self.rotation = np.moveaxis(m, (0, 1), (-2,-1))


# set default rotator
Rotator = QERRotator

def rotate(v, w, time = None, phase = None):
    """
    just do rotation of v about w
    """
    if phase is not None and time is None:
        time = 0
    return Rotator(w, time = time, phase = phase)(v)

def f2q(f):
    """
    convert array shape + (4,) to quaternion
    """
    return quat.from_float_array(np.asarray(f))

def q2f(q):
    """
    convert quaternion to array shape + (4,)
    """
    return quat.as_float_array(np.asarray(q))

def v2f(v):
    """
    covert vector to array shape + (4,) with [..., 0] set to 0
    """
    v = np.asarray(v)
    f = np.empty(v.shape[:-1] + (4,))
    f[..., 0 ] = 0.
    f[..., 1:] = v
    return  f

def v2q(v):
    """
    convert vector to verctorial part of quaternion
    """
    return  quat.from_float_array(v2f(v))

def f2v(f):
    """
    convert array shape + (4,) to vector (return just [..., 1:])
    """
    return  np.asarray(f)[..., 1:]

def q2v(q):
    """
    covert quaternion to vector (return just [..., 1:])
    """
    return  f2v(quat.as_float_array(np.asarray(q)))

def f2sv(f):
    """
    convert array shape + (4,) to scalar [..., 0] and vector [..., 1:] parts
    """
    f = np.asarray(f)
    return f[..., 0], f[..., 1:]

def q2sv(q):
    """
    convert quaternion to scalar [..., 0] and vector [..., 1:] parts
    """
    return  f2sv(quat.as_float_array(np.asarray(q)))

def sv2f(s, v):
    """
    convert scalar [...] and vector [..., 1:] arrays to shape + (4,)
    """
    s = np.asarray(s)
    v = np.asarray(v)
    f = np.empty(s.shape + (4,))
    f[..., 0 ] = s
    f[..., 1:] = v
    return  f

def sv2q(*args):
    """
    convert scalar [...] and vector [..., 1:] arrays to quaternion
    """
    return quat.from_float_array(sv2f(*args))

def w2f_(w):
    """
    convert rotation vector to quaternion in array shape + (4,)
    """
    w = np.asarray(w)
    f = np.empty(w.shape[:-1] + (4,))
    r = np.maximum(LA.norm(w, axis=-1), 1.e-99)
    rh = r * 0.5
    y = np.sin(rh) / r
    f[..., 0] = np.cos(rh)
    f[..., 1:] = w[..., :] * y[..., np.newaxis]
    return f

def w2q_(w):
    """
    convert rotation vector to quaternion
    """
    return quat.from_float_array(w2f_(w))

def w2q(w):
    """
    convert rotation vector to quaternion
    """
    w = np.asarray(w)
    f = np.zeros(w.shape[:-1] + (4,))
    f[..., 1:] = w[..., :] * 0.5
    return np.exp(quat.from_float_array(f))

def w2f(w):
    """
    convert rotation vector to quaternion in array shape + (,4)
    """
    return quat.as_float_array(w2q(w))

def f2w_(f):
    """
    convert quaternion as shape + (4,) to rotation vector
    """
    f = np.asarray(f)
    w = np.empty(f.shape[:-1] + (3,))
    r = np.arccos(f[..., 0]) * 2.
    sir = np.full_like(r, 2.)
    if np.shape(r) == ():
        if np.abs(r) > 1.e-99:
            sir[()] = r[()] / np.sin(r[()] * 0.5)
    else:
        ii = np.abs(r) > 1.e-99
        sir[ii] = r[ii] / np.sin(r[ii] * 0.5)
    w[..., :] = f[..., 1:] * sir[..., np.newaxis]
    return w

def q2w_(q):
    """
    convert quaternion to to rotation vector
    """
    return f2w_(quat.as_float_array(np.asarray(q)))

def f2w__(f):
    """
    convert quaternion as shape + (4,) to rotation vector

    This version seems slower but I should assume it be more precise.
    """
    f = np.asarray(f)
    n = LA.norm(f[..., 1:], axis = -1)
    if np.shape(n) == ():
        if np.abs(n) > 1.e-99:
            x = 2. * np.arctan2(n, f[0]) / n
        else:
            x = 2. / f[0]
    else:
        ii = np.abs(n) > 1.e-99
        x = np.empty_like(n)
        x[ii] = 2. * np.arctan2(n[ii], f[..., 0][ii]) / n
        ii = ~ii
        x[ii] = 2. / (f[..., 0][ii])
    return f[..., 1:] * x[..., np.newaxis]

def q2w__(q):
    """
    convert quaternion to to rotation vector
    """
    return f2w__(quat.as_float_array(np.asarray(q)))


def q2w(q):
    """
    convert quaternion to to rotation vector
    """
    return quat.as_float_array(np.log(np.asarray(q)) * 2.)[..., 1:]

def f2w(f):
    """
    convert quaternion as shape + (4,) to rotation vector
    """
    return q2w(quat.from_float_array(np.asarray(f)))


# ----------------------------------------------------------------------
# spinor matrix representation
# ----------------------------------------------------------------------
#
# v = (x, y, z), w = 0
# r = (w, x, y, z)
#
# s = [[ z + w  ,  x - y i],
#      [ x + y i, -z + w  ]]

sigma = np.asarray([
    np.array([[1 , 0 ], [ 0 , 1 ]], dtype = np.complex128),
    np.array([[0 , 1 ], [ 1 , 0 ]], dtype = np.complex128),
    np.array([[0, -1j], [ 1j, 0 ]], dtype = np.complex128),
    np.array([[1,  0 ], [ 0 , -1]], dtype = np.complex128),
           ])

def v2s(v):
     """
     transform vector (3,) to spinor
     """
     return np.tensordot(np.asarray(v), sigma[1:], axes = (-1,-3))

def f2s(f):
    """
    transform quaternion as (4,) to spinor
    """
    return  np.tensordot(np.asarray(f), sigma, axes = (-1,-3))

def s2v(s):
    """
    convert spinor to vector as (4,) (average duplicates)
    """
    s = np.asarray(s)
    return 0.5 * np.tensordot(s, sigma[1:].conj(), ((-1,-2),(-1,-2))).real

def s2f(s):
    """
    convert spinor to quaternion as (4,) (average duplicates)
    """
    s = np.asarray(s)
    return 0.5 * np.tensordot(s, sigma.conj(), ((-1,-2),(-1,-2))).real

def s2q(u):
    """
    convert spinor to quaternion (average duplicates)
    """
    return quat.from_float_array(s2f(s))

def q2s(q):
    """
    transform quaternion as (4,) to spinor
    """
    return f2s(q2f(q))

def w2s(w):
    """
    transform rotation vector to spinor
    """
    return f2s(w2f(w))

def s2w(s):
    """
    transform spinor to rotation vector
    """
    return f2w(s2f(w))


# ----------------------------------------------------------------------
# SU(2) matrix representation
# ----------------------------------------------------------------------
#
# q = a + b i + c j + d k
#
# u = [[ a + b i, c + d i],
#      [-c + d i, a - b i]]

M = np.asarray([
    np.array([[1,  0 ], [ 0 , 1 ]], dtype = np.complex128),
    np.array([[1j, 0 ], [ 0 ,-1j]], dtype = np.complex128),
    np.array([[0,  1 ], [-1 , 0 ]], dtype = np.complex128),
    np.array([[0,  1j], [ 1j, 0 ]], dtype = np.complex128),
           ])

def f2u(f):
    """
    transform quaternion as (4,) to SU(2)
    """
    return np.tensordot(np.asarray(f), M, axes = (-1,-3))

def q2u(q):
    """
    transform quaternion to SU(2)
    """
    return f2u(quat.as_float_array(np.asarray(q)))

def v2u(v):
    """
    transform vector as (3,) to SU(2)
    """
    f = np.zeros(v.shape[:-1] + (4,))
    f[...,1:] = v
    return f2u(f)

def w2u(w):
    """
    transform omega vector to SU(2)
    """
    return q2u(w2q(w))

def u2f(u):
    """
    convert SU(2) to quaternion as (4,) (average duplicates)
    """
    u = np.asarray(u)
    return 0.5 * np.tensordot(u, M.conj(), ((-1,-2),(-1,-2))).real

def u2q(u):
    """
    convert SU(2) to quaternion (average duplicates)
    """
    return quat.from_float_array(u2f(u))

def u2w(u):
    """
    convert SU(2) to rotation vector
    """
    return q2w(u2q(u))

def u2v(u):
    """
    convert SU(2) representation to vector of quaternion as (3,) shape (average duplicates)
    """
    f = u2f(u)
    return f[...,1:]

# ----------------------------------------------------------------------
# 4-Matrix representation
# ----------------------------------------------------------------------

_repM_ii = np.arange(3)
_repM_perm = np.array(list(permutations(_repM_ii + 2)))
_repMsig = list()
for i in range(8):
    ii = 1 - 2 * (np.right_shift(i, _repM_ii) % 2)
    _repMsig.append(_repM_perm * ii)
_repMsig = np.vstack(_repMsig)

def repMsig(i=0):
    assert 0 <= i < 48, "index needs to be in range [0, 48])"
    return _repMsig[i]

_repMtab = np.array(
    [[1,  2,  3,  4],
     [2, -1,  4, -3],
     [3, -4, -1,  2],
     [4,  3, -2, -1]])

def repMap(i=0):
    """
    return map for matrix representation
    """
    sig = repMsig(i)
    t = _repMtab.copy()
    t[1:,:] = +t[abs(sig)-1,:] * np.sign(sig)[:, None]
    t[:,1:] = -t[:,abs(sig)-1] * np.sign(sig)[None, :]
    return t

_repM_cache = dict()

def repM(i=0):
    """
    return list of matrixes for representation
    """
    if i in _repM_cache:
        return _repM_cache[i]
    t = repMap(i)
    mm = list()
    for i in range(4):
        ii = np.abs(t) == i + 1
        m = t.copy()
        m[ii] = np.sign(m[ii])
        m[~ii] = 0
        mm.append(m)
    _repM_cache[i] = mm
    return np.array(mm)

def f2m(f, i=0):
    """
    transform quaternion as (4,) to 4x4 matrix representaion
    """
    mm = repM(i)
    return np.tensordot(np.asarray(f), mm, axes = (-1,-3))

def v2m(v, i=0):
    """
    transform vector as (3,) to 4x4 matrix representaion
    """
    mm = repM(i)
    v = np.asarray(v)
    f = np.zeros(v.shape[:-1] + (4,))
    f[...,1:] = v
    return f2m(f)

def q2m(q, *args, **kwargs):
    """
    transform quaternion to 4x4 matrix representaion
    """
    return f2m(quat.as_float_array(np.asarray(q)), *args, **kwargs)

def checkM(m, i=0):
    """
    check whether m is 4x4 matrix representation of quaternion
    """
    mm = repM(i)
    c =  m[...,np.newaxis,:,:] * mm
    ii = [mx != 0 for mx in mm]
    c = c[...,ii]
    c = c.reshape(c.shape[:-1] + (4, 4,))
    return np.allclose(c, c[...,0:1])

def findMindex(m):
    """
    return all indices that match matrix m
    """
    finds = [i for i in range(48) if checkM(m, i)]
    if len(finds) == 0:
        return None
    if len(finds) == 1:
        return finds[0]
    return finds

def m2f(m, i=0):
    """
    convert 4x4 marix representation to quaternion as (4,) shape (average copies)
    """
    mm = repM(i)
    m = np.asarray(m)
    return 0.25 * np.tensordot(m, mm, ((-1,-2),(-1,-2)))

def m2v(*args, **kwargs):
    """
    convert 4x4 marix representation to vector of quaternion as (3,) shape (average copies)
    """
    f = m2f(*args, **kwargs)
    return f[...,1:]

def m2q(*args, **kwargs):
    """
    convert 4x4 marix representation to quaternion (average copies)
    """
    return quat.from_float_array(m2f(*args, **kwargs))

def w2m(w, i=0):
    """
    transform omega vector to 4x4 matrix representation
    """
    return q2m(w2q(w), i)

def m2w(m, i=0):
    """
    convert 4x4 matrix representaion to rotation vector
    """
    return q2w(m2q(m, i))

# ----------------------------------------------------------------------
# polar vectors
# ----------------------------------------------------------------------

def w2p(w):
    """
    convert rotation vector in xyz to polar (r,theta,phi)
    """
    w = np.asarray(w)
    p = np.empty(w.shape)
    p[..., 0] = LA.norm(w, axis = -1)
    p[..., 1] = np.arctan2(LA.norm(w[..., :2], axis = -1), w[..., 2])
    p[..., 2] = np.arctan2(w[..., 1], w[..., 0])
    return p

def p2w(p):
    """
    convert rotation vector in polar (r,theta,phi) to xyz
    """
    p = np.asarray(p)
    w = np.empty(p.shape)
    z12 = np.exp(p[..., 1:] * 1j)
    xy = z12[..., 0].imag * p[..., 0]
    w[..., 0] = xy * z12[..., 1].real
    w[..., 1] = xy * z12[..., 1].imag
    w[..., 2] = z12[..., 0].real * p[..., 0]
    return w

def l2p(latitude, longitude=None, radius=None):
    """
    convert longitude, lattitude[, radius = 1] to polar
    """
    if longitude is None:
        assert latitude.shape[-1] in (2, 3,)
        longitude = latitude[..., 1]
        if latitude.shape[-1] == 3:
            assert radius is None
            radius = latitude[..., 2]
        latitude = latitude[...,0]
    if radius is None:
        radius = 1.
    if np.shape(radius) == () and np.shape(latitude) != ():
        radius = np.tile(radius, longitude.shape)
    p = np.array([radius, (90 - latitude) * deg2rad, longitude * deg2rad])
    if p.ndim > 1:
        p = np.moveaxis(p, 0, -1)
    return p

def p2l(p, return_radius=False, return_vector=False):
    """
    convert polar to logitude, latitude
    """
    radius = p[..., 0]
    latitude  = 90 - p[..., 1] * rad2deg
    longitude = p[..., 2] * rad2deg
    if return_radius:
        v =  (latitude, longitude, radius,)
    else:
        v = (latitude, longitude,)
    if return_vector:
        return np.moveaxis(np.array(v), 0, -1)
    return v

def rotatel(*args, return_radius=None, return_vector=None, return_split=None):
    """
    rotate coordinates in (latititude, longitude[, radius]) about w

    can be individual vectors, scalars, or last dimension (one argument)

    call signatures:
    latlongrad,  w
    latlong, [r], w
    lat, long, [r], w

    detection for is not perfect when here is three positional
    arguments, there is degenerate cases, so please specify desired
    output.

    """
    assert 2 <= len(args) <=4
    return_split_ = False
    if len(args) == 4:
        return_radius_ = True
        return_vector_ = False
    elif len(args) == 2:
        return_vector_ = True
        return_radius_ = np.shape(args[0])[-1] == 3
    else:
        if np.shape(args[0]) == np.shape(args[1]):
            return_radius_ = False
            return_vector_ = False
        else:
            assert np.shape(args[0]) == 2
            return_radius_ = True
            return_vector_ = True
            return_split_ = True
    if return_split is None:
        return_split = return_split_
    if return_radius is None:
        return_radius = return_radius_
    if return_vector is None:
        return_vector = return_vector_

    w = args[-1]
    p = l2p(*args[:-1])
    v = p2w(p)
    v = rotate(v, w)
    p = w2p(v)
    l = p2l(p, return_radius=return_radius, return_vector=return_vector)
    if return_split:
        return l[...,:2], l[...,2]
    return l


def w2c(w):
    """
    convert rotation vector in xyz to cylindrical (rho,phi,z)
    """
    w = np.asarray(w)
    c = np.empty(w.shape)
    c[..., 0] = LA.norm(w[..., :2], axis = -1)
    c[..., 1] = np.arctan2(w[..., 1], w[..., 0])
    c[..., 2] = w[..., 2]
    return c

def c2w(c):
    """
    convert rotation vector in cylindrical (rho,phi,z) to xyz
    """
    c = np.asarray(c)
    w = np.empty(c.shape)
    w[..., 0] = c[..., 0] * np.cos(c[..., 1])
    w[..., 1] = c[..., 0] * np.sin(c[..., 1])
    w[..., 2] = c[..., 2]
    return w


def rotate2(w, w0 = None, rotate = None, rotate0 = None, align = None, lim = 1.e-99):
    """
    rotation vector to rotate from w0 to w.

    Default for w0 = unit vector in x-direction

    For anti-aligned vectors (sin(x) < 1e-99) the rotation axis is chosen
    sort of by chance.

    We could also just use arccos of dot product, may be faster though
    less accurate.

    Need vectorised version, but cases are a pain.

    rotate0 - rotation about initial axis applied first

    rotate - rotation about final axis applied last

    align - keep rotation vector on smae hemisphere as align, flip if needed.
    """
    w = np.asarray(w)
    l = np.asarray(LA.norm(w, axis = -1))
    l[l==0] = 1.
    w = w / l[..., np.newaxis]
    if w0 is None:
        w0 = np.tile(ex, w.shape[:-1] + (1,))
    else:
        w0 = np.asarray(w0)
        l = np.asarray(LA.norm(w0, axis = -1))
        l[l==0] = 1.
        w0 = w0 / l[..., np.newaxis]
        if w0.shape == (3,):
            w0 = np.tile(w0, w.shape[:-1] + (1,))
        if (w0.shape != w.shape):
            raise AttributeError(f'[rotate2] dimension mismatch {w.shape=}, {w0.shape=}')
    if align is not None:
        l = np.asarray(LA.norm(align, axis = -1))
        l[l==0] = 1.
        align = align / l[..., np.newaxis]

    x = np.cross(w0, w)
    s = np.asarray(LA.norm(x, axis=-1))
    c = np.asarray(np.sum(w0 * w, axis = -1))

    p = np.asarray(np.arctan2(s, c))
    if align is not None:
        ii = np.sum(x * align, axis=-1) < 0
        p[ii] -= 2 * np.pi
    s[s==0] = 1.
    w1 = x * (p / s)[..., np.newaxis]
    if rotate0 is not None:
        r = w0 * rotate0
        w1 = wchain(r, w1)
    if rotate is not None:
        r = w * rotate
        w1 = wchain(w1, r)
    return w1[()]


def rotate3(w, w1, w2):
    """
    get rotator for rotation about w from w1 to w2
    """
    w = np.asarray(w)
    w1 = np.asarray(w1)
    w2 = np.asarray(w2)

    # compute normal componnents
    p1 = np.cross(w, w1, axis = -1)
    p2 = np.cross(w, w2, axis = -1)

    wx = rotate2(p2, p1)

    align = np.sum(w * wx, axis = -1)
    if align < 0:
        wx *= 1 - np.pi * 2 / LA.norm(wx)
    return wx


def rotstep(w, v, steps = 0.5):
    """
    rotate v about w in steps
    """
    if np.shape(steps) != ():
        f2 = np.asarray(steps)
    elif steps <= 1:
        f2 = np.asarray(steps)
    else:
        f2 = np.arange(steps) / (steps - 1)
    return  Rotator(w, time = f2)(v)


def rotstep3(w, w1, w2, steps = 0.5):
    """
    rotate about w in steps from w1 to w2

    sort of like precession
    """
    w = np.asarray(w)
    w1 = np.asarray(w1)
    w2 = np.asarray(w2)

    wx = Rotate3(w, w1, w2)

    n1 = LA.norm(w1, axis = -1)
    n2 = LA.norm(w2, axis = -1)

    if np.shape(steps) != ():
        f2 = np.asarray(steps)
    elif steps <= 1:
        f2 = np.asarray(steps)
    else:
        f2 = np.arange(steps) / (steps - 1)
    f1 = np.asarray(1 - f2)

    return  Rotator(wx, time = wx)(w1 / n1) * \
        np.asarray(f1 * n1 + f2 * n2)[..., np.newaxis]

def wv2tr(w, v):
    """
    Decompose rotation of v about w into twist and rotation
    components, in that order.

    When recompositing, twist needs to be applied first.

    This routine seems highly inefficient.
    """
    w = np.asarray(w)
    v = np.asarray(v)
    v1 = Rotator(w)(v)
    r = rotate2(v1, v)
    t = wchain(w, -r)
    return t, r

def wv2rt(w, v):
    """
    Decompose rotation of v about w into rotation and twist
    components, in that order.

    When recompositing, rotation needs to be applied first.

    This routine seems highly inefficient.
    """
    w = np.asarray(w)
    v = np.asarray(v)
    v1 = Rotator(w)(v)
    r = rotate2(v1, v)
    t = wchain(-r, w)
    return r, t

# TODO - add routine(s) that computes projected twist (pendulum,
# coriolis) for rotation of v about w.  See sphere.items.Arrow
# Could also compute combined rotation vector, or variants using
# rotate2 (geodesic) decompositon.  Needs tests.

def wchain(w1, w2):
    """
    Rotation vector resulting from applying rotation w1 than w2.

    Not commutative.  Done by multiplying associated quaternions
    """
    return q2w(w2q(w2) * w2q(w1))

def wdiff(w1, w2):
    """
    Rotation vector that when applied after rotation w1 gives rotation w2
    """
    return wchain(-w1, w2)

# https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation
def wchain_(w1, w2):
    """
    Rotation composition using Rodrigues Formula.

    For testing only, do not use that.
    """
    a = np.asarray(w1)
    b = np.asarray(w2)
    an = LA.norm(a, axis = -1)
    bn = LA.norm(b, axis = -1)
    ae = a / an
    be = b / bn
    a2 = an * 0.5
    b2 = bn * 0.5
    ab = np.sum(ae * be, axis = -1)
    bxa = np.cross(be, ae, axis = -1)
    g2 = np.arccos(np.cos(a2) * np.cos(b2) - np.sin(a2) * np.sin(b2) * ab)
    ta2 = np.tan(a2)
    tb2 = np.tan(b2)
    tg2 = np.tan(g2)
    return (
        ae * ta2[..., np.newaxis] +
        be * tb2[..., np.newaxis] +
        bxa * (ta2 * tb2)[..., np.newaxis]) * (
            2. * g2 / (
            (1. - ta2 * tb2 * ab) * tg2))[..., np.newaxis]

# https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation
def wchain__(w1, w2):
    """
    Rotation composition using Rodrigues Formula.

    For testing only, do not use that.
    """
    a = np.asarray(w1)
    b = np.asarray(w2)
    an = norm(a)
    bn = norm(b)
    a2 = an * 0.5
    b2 = bn * 0.5
    ae = a / an
    be = b / bn
    bxa = cross(be, ae)
    sa2 = np.sin(a2)
    ca2 = np.cos(a2)
    sb2 = np.sin(b2)
    cb2 = np.cos(b2)
    ab = dot(ae, be)
    g2 = np.arccos(ca2 * cb2 - sa2 * sb2 * ab)
    d = sb2*ca2*be + sa2*cb2*ae + sb2*sa2*bxa
    return 2 * d * g2 / np.sin(g2)

def rotscale(w1, w2, steps = 0.5, rotation = None, align = None):
    """
    Smoothly rotate and scale w1 to w2

    If `steps` is an array return vectors for these fractions.

    If `steps` > 1 and integer type, compute that many steps including 0 and 1.

    If `steps` <= 1 or float scalar, computer that single vector.

    NOTE - It is most safe, and recommended, to use array notation.
           `steps=2` and `steps=2.` give different results.

    if `rotation` is not None, perform gradual rotation according to
    step, gradually scaling rotation angle from w1 to w2.  This
    internal rotation is applied after the rotscale operation about the
    final rotation axis.

    TODO - milti-D and contractions w/r w1, w2, steps (einsum)
    """
    steps = np.asarray(steps)
    if steps.ndim != 0:
        f2 = steps
    elif steps > 1 and np.issubsctype(steps, np.int64):
        f2 = np.arange(steps) / (steps - 1)
    else:
        f2 = steps

    f1 = np.asarray(1 - f2)
    n1 = LA.norm(w1, axis = -1)
    n2 = LA.norm(w2, axis = -1)
    w =  Rotator(np.multiply.outer(f2, rotate2(w2, w1, align=align)))(w1 / n1) * \
        np.asarray(f1 * n1 + f2 * n2)[..., np.newaxis]
    if rotation is not None:
        r = w  * (rotation * f2 / LA.norm(w, axis = -1))[..., np.newaxis]
        w = wchain(w, r)
    return w

def w2axyz(w, deg = False, lim = 1.e-12):
    """
    get rotation angle and xyz vector of axis (for use with vtk)
    """
    w = np.asarray(w, dtype = np.float64)
    a = LA.norm(w, axis = -1)
    ii = a < lim
    if np.shape(ii) == ():
        if ii:
            w = np.array([1., 0., 0.,])
            a = 0
        else:
            w = w / a
    else:
        jj = ~ii
        w = w.copy()
        w[ii, :] = np.array([1., 0., 0.])
        a[ii] = 0
        w[jj, :] /= a[jj, np.newaxis]

    if deg:
        a *= 180 / np.pi
    axyz = np.empty(np.shape(a) + (4,))
    axyz[..., 0] = a
    axyz[..., 1:] = w
    return axyz

#######################################################################
# [CHECK] check order of matrix
#######################################################################

def q2A(q):
    """
    convert quaternion to rotation matrix
    """
    f = quat.as_float_array(np.asarray(q))
    a, b, c, d = np.moveaxis(f, -1, 0)
    aa, bb, cc, dd = a * a, b * b, c * c, d * d
    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
    A = np.array(
        [[aa + bb - cc - dd, 2 * (bc - ad), 2 * (bd + ac)],
         [2 * (bc + ad), aa + cc - bb - dd, 2 * (cd - ab)],
         [2 * (bd - ac), 2 * (cd + ab), aa + dd - bb - cc]])
    ii = np.arange(len(f.shape) + 1)
    jj = np.roll(ii, 2)
    return np.moveaxis(A, ii, jj)

def w2A(w):
    """
    convert rotation vector to rotation matrix
    """
    return q2A(w2q(w))

def f2A(f):
    """
    convert quaternion as shape + (4,) to rotation matrix
    """
    return q2A(quat.from_float_array(np.asarray(f)))

def A2f(A):
    """
    convert matrix to quaternion as M shape + (4,)
    """
    A = np.asarray(A)
    f = np.empty(A.shape[:-2] + (4,))
    ii = np.array([0,1,2])
    m00, m11, m22 = A[..., ii, ii]
    f[..., 0] = math.sqrt(np.maximum( 0, 1 + m00 + m11 + m22 ) ) * 0.5
    f[..., 1] = math.sqrt(np.maximum( 0, 1 + m00 - m11 - m22 ) ) * 0.5
    f[..., 2] = math.sqrt(np.maximum( 0, 1 - m00 + m11 - m22 ) ) * 0.5
    f[..., 3] = math.sqrt(np.maximum( 0, 1 - m00 - m11 + m22 ) ) * 0.5
    f[..., 1] *= np.sign(A[..., 2, 1] - A[..., 1, 2])
    f[..., 2] *= np.sign(A[..., 0, 2] - A[..., 2, 0])
    f[..., 3] *= np.sign(A[..., 1, 0] - A[..., 0, 1])
    return f

def A2q(A):
    """
    convert matrix to quaternion
    """
    return quat.from_float_array(A2f(A))

def A2w(A):
    """
    convert matrix to rotation vector
    """
    return q2w(A2q(A))

def angle(w, deg=False):
    """
    return rotation angle of vector
    """
    a = LA.norm(w)
    if deg:
        a *= rad2deg
    return a[()]

def angle2(w1, w2, deg=False, line=False):
    """
    compute angle between 2 vertors or vector arrays

    dimensions need to be broadcastable

    `line` - return angle of intersection of lines raher than angle
    between vectors.  The result is always in Q1.
    """
    w1 = np.asarray(w1, dtype=np.float64)
    w2 = np.asarray(w2, dtype=np.float64)
    s = norm(cross(w1, w2))
    c = dot(w1, w2)
    if line:
        c = np.abs(c)
    a = np.arctan2(s, c)
    if deg:
        a *= rad2deg
    return a[()]

def angle3(w, w1, w2, deg=False, positive=True):
    """
    get rotation angle for rotation about w from w1 to w2
    """
    w = np.asarray(w, dtype=float)
    w1 = np.asarray(w1, dtype=float)
    w2 = np.asarray(w2, dtype=float)

    p1 = np.cross(w, w1, axis=-1)
    p2 = np.cross(w, w2, axis=-1)

    n = np.asarray(LA.norm(w, axis=-1))
    n[n==0] = 1.

    c = np.sum(p1 * p2, axis=-1) * n
    s = np.sum(np.cross(p1, p2, axis=-1) * w, axis=-1)
    a = np.asarray(np.arctan2(s, c))
    if positive:
        a[a < 0] += np.pi * 2
    if deg:
        a *= rad2deg
    return a[()]


def random_unit_vector(*dims):
    """
    return an array of random unit vectors
    """
    if len(dims) == 1 and isinstance(dims[0], tuple):
        dims = dims[0]
    z, p = np.random.random((2,) + dims)
    z = 2 * z - 1
    xy = np.sqrt(1-z**2) * np.exp(2j * np.pi * p)
    w = np.empty(dims + (3,))
    w[..., 0] = xy.real
    w[..., 1] = xy.imag
    w[..., 2] = z
    return w


def random_rotation_vector(*dims):
    """
    return an array of random rotation vectors
    """
    if len(dims) == 1 and isinstance(dims[0], tuple):
        dims = dims[0]
    z, p, a = np.random.random((3,) + dims)
    z = 2 * z - 1
    xy = np.sqrt(1-z**2) * np.exp(2j * np.pi * p)
    w = np.empty(dim + (3,))
    w[..., 0] = xy.real
    w[..., 1] = xy.imag
    w[..., 2] = z
    return w * (np.pi * a)


def dotcross(a, b):
    """
    return dot product and cross product of two vectors (using quaternions)
    """
    d, c = q2sv(v2q(a) * v2q(b))
    return -d, c


def v_dot_f(v, f):
    """
    vector-quaternion product (vector form)
    """
    return q2f(v2q(v) * f2q(f))


def f_dot_f(f1, f2):
    """
    quaternion product for vector form
    """
    return q2f(f2q(f1) * f2q(f2))


def conj_f_(f):
    """
    conjugate in quaternion vector form

    this versions seems, unexpectely, 2-3x slower
    """
    f1 = np.empty_like(f)
    f1[..., 0 ] =  f[..., 0 ]
    f1[..., 1:] = -f[..., 1:]
    return f1


def conj_f(f):
    """
    conjugate in quaternion vector form
    """
    return q2f(f2q(f).conj())