"""
Utility routines for rotations and vectors.

Default is to assume arrayc of vectors with vector component being
last.

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

import numpy as np
from numpy import linalg as LA

from .const import *

iv = (..., None,)

def normalize(v, axis=-1):

    """
    normalize vector  (default: along last axis)
    """
    v = np.asarray(v)
    n = np.expand_dims(LA.norm(v, axis=axis), axis=axis)
    n[n==0] = 1.
    return v / n

def cross(v, w, axis=-1):
    """
    return cross product (default: along last axis)
    """
    v = np.asarray(v)
    w = np.asarray(w)
    return np.cross(v, w, axis=axis)

def normcross(v, w, axis=-1):
    """
    return norm of cross product (default: along last axis)
    """
    return norm(cross(v, w, axis=axis), axis=axis)

def cross_vec(v, w, axis=-1, axisa=None, axisb=None, axisc=None):
    """
    intended to try whether multiplications vectorise but does not
    seem to be the case.

    """
    if axisa is None:
        axisa = axis
    if axisb is None:
        axisb = axis
    if axisc is None:
        axisc = axis
    v = np.moveaxis(np.asarray(v), axisa, 0)
    w = np.moveaxis(np.asarray(w), axisb, 0)
    return np.reshape(np.array([
        v[1]*w[2] - v[2]*w[1],
        v[2]*w[0] - v[0]*w[2],
        v[0]*w[1] - v[1]*w[0],
         ]), 0, axisc)

def outer_cross(a, b, axis=-1, axisa=None, axisb=None, axisc=None, align=0):
    """
    outer cross product of multi-D arrays

    The last non-verctor align dimensions of a and the first align
    dimensions of b are 'aligned', i.e., just vectorized and need to
    be same.

    TODO - allow generic align as tuple of list of dimensions
    """
    if axisa is None:
        axisa = axis
    if axisb is None:
        axisb = axis
    if axisc is None:
        axisc = axis

    a = np.moveaxis(np.asarray(a), axisa, -1)
    b = np.moveaxis(np.asarray(b), axisb, -1)

    if align > 0:
        assert a.shape[-align:] == b.shape[:align]

    da = a.shape[:-1]
    db = b.shape[:-1]
    na = len(da)
    nb = len(db)
    sa = da + (1,) * (nb - align) + (3,)
    sb = (1,) * (na - align) + db + (3,)

    a = v.reshape(sa)
    b = v.reshape(sb)

    return np.reshape(np.cross(a, b, axis=-1), -1, axisc)


def dot(v, w, axis=-1):
    """
    return dot product (default: along last axis)
    """
    v = np.asarray(v)
    w = np.asarray(w)
    return np.sum(v * w, axis=-1)

def norm(v, axis=-1):
    """
    return norm of vector v, |v|
    """
    return np.sqrt(dot(v, v, axis=axis))

def norme(v, axis=-1):
    """
    return norm and unit vector of vector v, |v|, v / |v|
    """
    v = np.array(v, dtype=np.float64)
    n = np.sqrt(dot(v, v, axis=axis))
    if n.shape == ():
        if n == 0:
            v[:] = 0.
        else:
            v = v / n
        return n, v
    ii = n==0
    n[ii] = 0.
    v[ii] = 0.
    ii = ~ii
    v[ii] = v[ii] * np.reciprocal(n[ii, np.newaxis])
    return n, v

def normr(v, axis=-1):
    """
    return reciprocal of norm of vector v, 1/|v|
    """
    return np.reciprocal(np.sqrt(dot(v, v, axis=axis)))

def norm2(v, axis=-1):
    """
    return square of norm of vector v, |v|**2
    """
    return dot(v, v, axis=axis)

def log_norm(v, axis=-1):
    """
    return log10 of norm of vector v, |v|**2
    """
    return np.log10(dot(v, v, axis=axis)) * 0.5

def normr2(v, axis=-1):
    """
    return reciprocal of square of norm of vector v, 1/|v|**2
    """
    return np.reciprocal(dot(v, v, axis=axis))

def cross21(a, b, a2 = None, axis=-1):
    """
    return a x (a x b) = a (a b) - b (a a)

    the negative of the componnent of b perpendicular to a
    (scaled by |a|**2)
    """
    if a2 is None:
        a2 = np.sum(a**2, axis=axis)
    return (a * np.expand_dims(np.sum(a * b, axis=axis), axis=axis)
            - b * np.expand_dims(a2, axis=axis))

def perpendicular1D(v):
    """
    find some perpendicular vector (1D version)
    """
    v = np.asarray(v)
    assert v.shape == (3,)
    i = np.argsort(np.abs(v))[0]
    n = basis[i].copy()
    if v[i] == 0:
        return n
    v = np.cross(n, v)
    return normalize(v)

def perpendicular(v, axis=-1):
    """
    find some perpendicular vector (multi-D version)
    """
    v = np.asarray(v)
    v = np.moveaxis(v, axis, -1)
    ii = np.argsort(np.abs(v), axis=-1)
    n = basis[ii[..., 0]].copy()
    w = np.take_along_axis(v, ii, axis=-1)
    jj = w[..., 0] != 0
    n[jj] = np.cross(n[jj], v[jj], axis=-1)
    n[jj] = normalize(n[jj])
    return np.moveaxis(n, -1, axis)