
import itertools, numpy as np, math
import uuid

from . import VectorOps as vec_ops
from . import SetOps as sets
from .Misc import is_numeric, is_zero

__all__ = [
    "permutation_sign",
    "levi_cevita_maps",
    "levi_cevita_tensor",
    "levi_cevita3",
    "levi_cevita_dot",
    "nca_op_deriv",
    "tensordot_deriv",
    "tensorprod_deriv",
    "scalarprod_deriv",
    "inverse_transformation",
    "optimizing_transformation",
    "matinv_deriv",
    "matdet_deriv",
    "matsqrt_deriv",
    "scalarinv_deriv",
    "scalarpow_deriv",
    "tensor_reexpand",
    "tensorops_deriv",
    "vec_norm_unit_deriv",
    "vec_angle_deriv",
    "vec_cross_deriv",
    "vec_anglecos_deriv",
    "vec_anglesin_deriv",
    "vec_dihed_deriv",
    "vec_plane_angle_deriv"
]

def permutation_sign(perm, check=True):
    # essentially a swap sort on perm
    # https://stackoverflow.com/a/73511014
    parity = 1
    perm = np.asanyarray(perm)
    if check:
        perm = np.argsort(np.argsort(perm))
    else:
        perm = perm.copy()
    for i in range(len(perm)):
        while perm[i] != i: # ensure subblock is sorted
            parity *= -1
            j = perm[i]
            perm[i], perm[j] = perm[j], perm[i]
    return parity
def levi_cevita_maps(k):
    perms = sets.permutation_indices(k, k)
    signs = np.array([permutation_sign(p, check=False) for p in perms])
    return perms, signs
def levi_cevita_tensor(k, sparse=False):
    pos, vals = levi_cevita_maps(k)
    if sparse:
        from .Sparse import SparseArray
        a = SparseArray.from_data(
            (
                pos.T,
                vals
            ),
            shape=(k,)*k
        )
    else:
        a = np.zeros((k,)*k, dtype=int)
        a[tuple(pos.T)] = vals
    return a
# levi_cevita3 = levi_cevita_tensor(3)
levi_cevita3 = np.array([
    [[0, 0, 0],
     [0, 0, 1],
     [0, -1, 0]],

    [[0, 0, -1],
     [0, 0, 0],
     [1, 0, 0]],

    [[0, 1, 0],
     [-1, 0, 0],
     [0, 0, 0]]
])

def sylvester_solve(A, B, C):
    A = np.asanyarray(A)
    B = np.asanyarray(B)
    C = np.asanyarray(C)
    n = A.shape[1]
    m = B.shape[0]
    S = np.kron(B.T, np.eye(n)) + np.kron(np.eye(m), A)
    val = np.linalg.solve(S, C.flatten()) # in case it can be cleverer than I can
    return val.reshape(m,n)

def levi_cevita_dot(k, a, /, axes, shared=None):
    pos, vals = levi_cevita_maps(k)
    return vec_ops.semisparse_tensordot((tuple(pos.T), vals, (k,) * k), a, axes, shared=shared)
    # a = np.asanyarray(a)
    # lv_axes = axes[0]
    # a_axes = axes[1]
    # if util.is_numeric(lv_axes):
    #     lv_axes = [lv_axes]
    # if util.is_numeric(a_axes):
    #     a_axes = [a_axes]
    # for ax in reversed(a_axes):
    #     a = np.moveaxis(a, ax, 0)
    #
    # nax = len(lv_axes)
    # if nax == k:
    #     new = np.dot(vals, a[tuple(pos.T)])
    # else:
    #     new_shape = (k,) * (k-nax) + a.shape[nax:]
    #     new = np.zeros(new_shape)
    #     idx = np.setdiff1d(np.arange(k), lv_axes)
    #     for p, v in zip(pos, vals):
    #         new[p[lv_axes]] += v*a[p[idx]]
    # return new

# levi_cevita3.__name__ = "levi_cevita3"
# levi_cevita3.__doc__ = """
#     The 3D Levi-Cevita tensor.
#     Used to turn cross products into matmuls
#     """



def get_nca_shifts(order, k):
    permute_pos = np.arange(k)
    ncombs = math.comb(order, k)
    shifts = np.broadcast_to(np.arange(order)[np.newaxis], (ncombs, order)).copy()
    for i,pos in enumerate(itertools.combinations(range(order), r=k)):
        shifts[i, pos] = permute_pos
    return shifts
def apply_nca_op(op, order, k, A_expansion, B_expansion, deriv_axis, a, b, contract, shared, identical):
    s = order - k
    if s >= len(A_expansion) or k >= len(B_expansion):
        return 0

    A = A_expansion[s]
    B = B_expansion[k]
    for T in [A, B]:
        if (
            is_zero(T)
            or (T.shape == () and T == 0)
        ):
            return 0
    if shared is None:
        shared = 0
    if contract: # axes disappear, so we just account for the shifts
        axes = [[x+s for x in a], [x+k for x in b]]
    else: # axes appeared, so we need to include those in the product
          # we _forced_ axes to be at the end of the arrays since it was too hard
          # to keep track of them otherwise...
          # actually I guess I could have put the derivative axes at the end...
          # and then the axes would never change...but that has other complications
        axes = [
            [shared + i for i in range(s)] + [x+s for x in a],
            [shared + i for i in range(k)] + [x+k for x in b]
        ]

    if shared == 0:
        base = op(A, B, axes=axes)
    else:
        base = op(A, B, axes=axes, shared=shared)

    # next we need to move all of the derivative axes in the second tensor to the beginning
    # so we can symmetrize
    d = deriv_axis + s
    for i in range(k):
        base = np.moveaxis(base, d+i, shared)
    part = [x for x in reversed(sorted([s,k])) if x > 0]
    if len(part) > 1:
        base = nca_symmetrize(base, part, shared=shared, identical=identical)
    return base
def nca_op_order_deriv(op, order, A_expansion, B_expansion, deriv_axis, a, b, contract, shared, identical):
    full = None
    for k in range(order+1):
        term = apply_nca_op(op, order, k, A_expansion, B_expansion, deriv_axis, a, b, contract, shared, identical)
        if full is None:
            full = term
        else:
            full = full + term
    return full
def nca_op_deriv(op,
                 A_expansion,
                 B_expansion,
                 order,
                 axes,
                 contract,
                 shared=None,
                 identical=False
                 ):
    A_expansion = [np.asanyarray(A) for A in A_expansion]
    B_expansion = [np.asanyarray(B) for B in B_expansion]

    a_ax, b_ax = axes
    if isinstance(a_ax, int): a_ax = [a_ax]
    if isinstance(b_ax, int): b_ax = [b_ax]
    a_ax = [ a if a >= 0 else A_expansion[0].ndim + a for a in a_ax ]
    b_ax = [ b if b >= 0 else B_expansion[0].ndim + b for b in b_ax ]

    if contract: # the derivative axis will always be at nA + 1 - num_contracted - the shared axes
        deriv_axis = A_expansion[0].ndim - len(a_ax)
    else:
        # we require that the outer product be ordered
        # so we now which axes to move around
        a_ax = np.sort(a_ax)
        a_dim = A_expansion[0].ndim
        if np.any(a_ax != np.arange(a_dim - len(a_ax), a_dim)):
            raise ValueError("axes {} must be the final axes of A".format(a_ax))
        b_ax = np.sort(b_ax)
        b_dim = B_expansion[0].ndim
        if np.any(b_ax != np.arange(b_dim - len(b_ax), b_dim)):
            raise ValueError("axes {} must be the final axes of B".format(b_ax))
        deriv_axis = a_dim

    if isinstance(order, int):
        order = list(range(1, order+1))

    # if shared is not None:
    #     deriv_axis = deriv_axis - shared

    derivs = [
        nca_op_order_deriv(op, o, A_expansion, B_expansion, deriv_axis, a_ax, b_ax,
                           contract, shared, identical
                           )
        for o in order
    ]

    return derivs

def _deriv_construct(zero_order_handler, expansion_generator, order):
    if is_numeric(order):
        order = list(range(order+1))
    high_order_expansion = expansion_generator([o for o in order if o > 0])

    # a lot of work to make sure we can do derivatives at a specifically targeted order
    n = 0
    final_expansion = []
    for o in order:
        if o == 0:
            final_expansion.append(
                zero_order_handler()
            )
        else:
            final_expansion.append(high_order_expansion[n])
            n += 1
    return final_expansion

def tensordot_deriv(A_expansion, B_expansion,
                    order,
                    axes=None,
                    shared=None,
                    identical=False
                    ):

    if len(A_expansion) == 0 or len(B_expansion) == 0:
        return [0]*((order+1) if is_numeric(order) else len(order))

    if axes is None: axes = [-1, 0]
    if shared is not None and shared > 0:
        base_op = vec_ops.vec_tensordot
    else:
        base_op = lambda a,b,axes=None,shared=None:np.tensordot(a, b, axes=axes)
    op = lambda a,b,axes=None, shared=None: (
        0
            if is_zero(a) or is_zero(b) else
        base_op(a, b, axes=axes, shared=shared)
    )
    return _deriv_construct(
        lambda : op(A_expansion[0], B_expansion[0], axes=axes, shared=shared),
        lambda ords: nca_op_deriv(op,
                     A_expansion, B_expansion,
                     ords,
                     axes=axes,
                     contract=True,
                     shared=shared,
                     identical=identical
                     ),
        order
    )

def prep_prod_arrays(A_expansion, a_ax):
    _ = []
    for n,a in enumerate(A_expansion):
        a = np.asanyarray(a)
        for x in a_ax:
            x = (a.ndim + x) if x < 0 else (x + n)
            a = np.moveaxis(a, x, -1)
        _.append(a)
    return _
def pre_broadcast_prod(A_expansion, B_expansion, axes):

    a_ax, b_ax = axes
    if isinstance(a_ax, int): a_ax = [a_ax]
    if isinstance(b_ax, int): b_ax = [b_ax]

    A_expansion = prep_prod_arrays(A_expansion, a_ax)
    B_expansion = prep_prod_arrays(B_expansion, b_ax)

    expand_dims = ...

def tensorprod_deriv(
        A_expansion, B_expansion,
        order,
        axes=None,
        identical=False
):
    if len(A_expansion) == 0 or len(B_expansion) == 0:
        return [0]*((order+1) if is_numeric(order) else len(order))
    if axes is None:
        axes = [-1, -1]
    elif isinstance(axes, str) and axes == 'all':
        axes = [np.arange(A_expansion[0].ndim), np.arange(B_expansion[0].ndim)]
    a_ax, b_ax = axes
    if isinstance(a_ax, int): a_ax = [a_ax]
    if isinstance(b_ax, int): b_ax = [b_ax]
    # a_ax, b_ax = pre_broadcast_prod(A_expansion, B_expansion, axes)

    shared = A_expansion[0].ndim - len(a_ax)
    # if is_numeric(A_expansion[0]) or is_numeric(B_expansion[0]):
    #     # scalar product
    #     return scalarprod_deriv(A_expansion, B_expansion, order, identical=identical)

    A_expansion = [np.asanyarray(a) for a in A_expansion]
    B_expansion = [np.asanyarray(b) for b in B_expansion]
    op = lambda left, right, axes=None, shared=None: (
        0
            if is_zero(left) or is_zero(right) else
        vec_ops.vec_outer(left, right, axes=axes, order=0)
    )
    return _deriv_construct(
        lambda: op(A_expansion[0], B_expansion[0], axes=axes, shared=shared),
        lambda ords: nca_op_deriv(op,
                                  A_expansion, B_expansion,
                                  ords,
                                  axes=[a_ax, b_ax],
                                  contract=False,
                                  shared=shared,
                                  identical=identical
                                  ),
        order
    )

def _scalar_prod(a, b):
    if is_numeric(a) or is_numeric(b):
        if is_zero(a) or is_zero(b):
            return 0
        else:
            return a * b
    else:
        return vec_ops.vec_outer(a, b, axes=[list(range(a.ndim)), list(range(b.ndim))], order=0)

def scalarprod_deriv(s_expansion, A_expansion,
                     order,
                     identical=False
                     ):
    s_expansion = [np.asarray(s) for s in s_expansion]
    A_expansion = [np.asarray(a) for a in A_expansion]
    s0_shape = s_expansion[0].shape if len(s_expansion) > 0 else ()
    A0_shape = A_expansion[0].shape if len(A_expansion) > 0 else ()
    shared = min([len(s0_shape), len(A0_shape)])
    return tensorprod_deriv(
        s_expansion,
        A_expansion,
        order=order,
        identical=identical,
        axes=[
            np.arange(shared, s_expansion[0].ndim),
            np.arange(shared, A_expansion[0].ndim)
        ]
    )
    if is_numeric(order):
        order = list(range(order+1))
    suborder = [o for o in order if o > 0]
    # s_expansion, A_expansion = pre_broadcast_scalar_mult(s_expansion, A_expansion)
    if is_numeric(A_expansion[0]):
        if is_numeric(s_expansion[0]):
            base_expansion = scalarprod_deriv(s_expansion, A_expansion[1:],
                                              [o - 1 for o in suborder],
                                              identical=identical)
        else:
            shared = min([s_expansion[0].ndim, A_expansion[0].ndim])
            base_expansion = tensorprod_deriv(s_expansion, A_expansion[1:],
                                              [o - 1 for o in suborder],
                                              identical=identical,
                                              axes=[
                                                  np.arange(shared, s_expansion[0].ndim),
                                                  np.arange(shared, A_expansion[0].ndim)
                                              ])
        rem_expansion = [
            _scalar_prod(s_expansion[o], A_expansion[0])
                if len(s_expansion) > o else 0
            for o in suborder
        ]
    elif not is_numeric(s_expansion[0]):
        shared = min([s_expansion[0].ndim, A_expansion[0].ndim])
        return tensorprod_deriv(s_expansion, A_expansion,
                                order,
                                identical=identical,
                                axes=[
                                    np.arange(shared, s_expansion[0].ndim),
                                    np.arange(shared, A_expansion[0].ndim)
                                ])
    else:
        base_expansion = tensorprod_deriv(s_expansion[1:], A_expansion,
                                          order=[o - 1 for o in suborder],
                                          axes=[
                                              np.arange(s_expansion[1].ndim),
                                              np.arange(A_expansion[0].ndim)
                                          ],
                                          identical=identical)
        rem_expansion = [
            _scalar_prod(s_expansion[0], A_expansion[o]) if len(A_expansion) > o else 0
            for o in suborder
        ]
    high_order_expansion = [x + y for x, y in zip(base_expansion, rem_expansion)]

    # a lot of work to make sure we can do derivatives at a specifically targeted order
    n = 0
    final_expansion = []
    for o in order:
        if o == 0:
            final_expansion.append(_scalar_prod(s_expansion[0], A_expansion[0]))
        else:
            final_expansion.append(high_order_expansion[n])
            n += 1
    return final_expansion

def inverse_transformation(forward_expansion, order, reverse_expansion=None, allow_pseudoinverse=False):
    if reverse_expansion is None:
        if allow_pseudoinverse and forward_expansion[0].shape[0] != forward_expansion[0].shape[1]:
            B = np.linalg.pinv(forward_expansion[0])
        else:
            B = np.linalg.inv(forward_expansion[0])
        reverse_expansion = [B]
    else:
        reverse_expansion = list(reverse_expansion)
        B = reverse_expansion[0]

    if not is_numeric(order): order = np.max(order)
    order = list(range(2, order+1))
    shared = B.ndim - 2

    for o in order:
        new_B = -tensor_reexpand(reverse_expansion, forward_expansion, [o])[-1]
        if isinstance(new_B, np.ndarray):
            # need to multiply in the inverse now, too
            new_B = vec_ops.vec_tensordot(new_B, B, axes=[-1, shared], shared=shared)
            # for _ in range(new_B.ndim - 1):
            #     new_B = np.tensordot(B, new_B, axes=[1, -2])
        reverse_expansion = reverse_expansion + [new_B]
    return reverse_expansion

def kron_prod_derivs(A_expansion, B_expansion, order):
    # s = A_expansion[0].ndim - 2
    C_expansion = tensorprod_deriv(A_expansion, B_expansion, order,
                                   axes=[[-1, -2], [-1, -2]]
                                   )
    # now we need to reshape based on the fact that we end up with nnxmm
    return [
        c.reshape(c.shape[:-4] + (c.shape[-4]*c.shape[-3], c.shape[-2]*c.shape[-1]))
        for c in C_expansion
    ]

def kron_sum_derivs(A_expansion, B_expansion, order):
    base_shape = B_expansion[0].shape[:-2]
    ns = len(base_shape)
    m = B_expansion[0].shape[-1]
    if ns > 0:
        I_expansion = np.expand_dims(np.eye(m), list(range(ns)))
        I_expansion = np.broadcast_to(I_expansion, base_shape + (m, m))
    else:
        I_expansion = np.eye(m)
    I_expansion = [I_expansion]
    left_expansion = kron_prod_derivs(A_expansion, I_expansion, order)

    n = A_expansion[0].shape[-1]
    if ns > 0:
        I_expansion = np.expand_dims(np.eye(n), list(range(ns)))
        I_expansion = np.broadcast_to(I_expansion, base_shape + (n, n))
    else:
        I_expansion = np.eye(n)
    I_expansion = [I_expansion]
    right_expansion = kron_prod_derivs(I_expansion, B_expansion, order)

    return [l+r for l,r in zip(left_expansion, right_expansion)]

def sylv_derivs(A_expansion, B_expansion, C_expansion, order, shared=None, ks_expansion=None, inv_expansion=None):
    n = A_expansion[0].shape[-1]
    m = B_expansion[0].shape[-1]
    B_expansion = [np.moveaxis(b, -1, -2) for b in B_expansion]
    if ks_expansion is not None:
        ks_expansion = list(ks_expansion) + kron_sum_derivs(A_expansion, B_expansion, order)
    else:
        ks_expansion = kron_sum_derivs(A_expansion, B_expansion, order)
    if isinstance(order, list):
        inv_expansion = matinv_deriv(ks_expansion, order[0], base_expansion=inv_expansion)
    else:
        inv_expansion = matinv_deriv(ks_expansion, order, base_expansion=inv_expansion)
    C_expansion = [np.reshape(c, c.shape[:-2] + (c.shape[-1] * c.shape[-2],)) for c in C_expansion]
    X_expansion = tensordot_deriv(C_expansion, inv_expansion, order,
                                  axes=[-1, -1],
                                  shared=shared
                                  )

    return [np.reshape(x, x.shape[:-1] + (n,m)) for x in X_expansion], (ks_expansion, inv_expansion)

def matsqrt_deriv(A_expansion, order):
    shared = A_expansion[0].ndim - 2
    if shared == 0: shared = None
    B_expansion = [
        vec_ops.frac_powh(A_expansion[0], 1/2, pow=lambda evals,_:np.sqrt(evals))
        # sqrt maybe faster than pow(..., 1/2)
    ]
    inv_expansion = None
    ks_expansion = None
    for o in range(1, order+1):
        new_expansion, (ks_expansion, inv_expansion) = sylv_derivs(B_expansion, B_expansion, A_expansion[1:],
                                                                   ks_expansion=ks_expansion,
                                                                   inv_expansion=inv_expansion,
                                                                   shared=shared,
                                                                   order=[o - 1])
        B_expansion = B_expansion + new_expansion

    return B_expansion

def matinv_deriv(A_expansion, order, base_expansion=None):
    shared = s = A_expansion[0].ndim - 2
    if shared == 0: shared = None
    if base_expansion is None:
        B = np.linalg.inv(A_expansion[0])
        base_expansion = [B]
    inverse_expansion = list(base_expansion)

    for o in range(len(inverse_expansion), order+1):
        expansion = tensorops_deriv(
            A_expansion[1:],
                [s, s + 1],
            inverse_expansion,
                [s, s + 0],
            inverse_expansion,
            order=[o - 1],
            shared=shared
        )
        inverse_expansion = inverse_expansion + [-expansion[-1]]
    return inverse_expansion

def matdet_deriv(forward_expansion, order):
    shared = forward_expansion[0].ndim - 2
    reverse_expansion = matinv_deriv(forward_expansion, order)
    tr_inner_exp = tensordot_deriv(forward_expansion[1:], reverse_expansion,
                                   order - 1,
                                   axes=[1+shared, 1+shared],
                                   shared=shared
                                   )
    tr_exp = [np.trace(a, axis1=-2, axis2=-1) for a in tr_inner_exp]
    det_exp = [np.linalg.det(forward_expansion[0])]
    for o in range(0, order):
        det_exp = det_exp + [scalarprod_deriv(det_exp, tr_exp, [o])[-1]]
    return det_exp

def _broadcast_mul(scalar, other):
    if isinstance(scalar, np.ndarray) and scalar.ndim < other.ndim:
        return np.expand_dims(scalar, list(range(scalar.ndim, other.ndim))) * other
    else:
        return scalar * other
def _scalarinv_deriv(scalar_expansion, o):
    from ..Combinatorics import IntegerPartitioner

    shared = scalar_expansion[0].ndim
    term = 0
    for parts in IntegerPartitioner.partitions(o):
        l = len(parts[0])
        scaling = ((-1) ** l) * math.factorial(l) / (scalar_expansion[0] ** (l + 1))
        term += sum(
            _broadcast_mul(scaling, nca_partition_prod(p, scalar_expansion[1:], shared=shared))
            for p in parts
        )
    return term

def scalarinv_deriv(scalar_expansion, order):
    return _deriv_construct(
        lambda : 1/scalar_expansion[0],
        lambda ords: [_scalarinv_deriv(scalar_expansion, o) for o in ords],
        order
    )

def _scalarpow_deriv(scalar_expansion, exp, o):
    from ..Combinatorics import IntegerPartitioner
    shared = scalar_expansion[0].ndim

    term = 0
    for parts in IntegerPartitioner.partitions(o):
        l = len(parts[0])
        factorial_terms = np.prod(exp - np.arange(l))
        if factorial_terms == 0:
            continue
        scaling = factorial_terms * (scalar_expansion[0] ** (exp - l))
        term += sum(
            _broadcast_mul(scaling, nca_partition_prod(p, scalar_expansion[1:], shared=shared))
            for p in parts
        )
    return term

def scalarpow_deriv(scalar_expansion, exp, order):
    return _deriv_construct(
        lambda : np.power(scalar_expansion[0], exp),
        lambda ords: [_scalarpow_deriv(scalar_expansion, exp, o) for o in ords],
        order
    )

def vec_norm_unit_deriv(vec_expansion, order, base_expansion=None):
    if base_expansion is None:
        a = vec_expansion[0]
        r = np.linalg.norm(a, axis=-1)
        if a.ndim > 1:
            u = a / r[..., np.newaxis]
        else:
            u = a / r
        base_expansion = [r, u]
    base_expansion = list(base_expansion)
    norm_inv_expansion = scalarinv_deriv(base_expansion, order=len(base_expansion)-2)

    a_expansion = [vec_expansion[0], vec_ops.identity_tensors(base_expansion[1].shape[:-1], 3)]

    if not is_numeric(order):
        order = max(order)

    for o in range(len(base_expansion), order+2):
        # if len(base_expansion) <= o:
        norm_inv_expansion = norm_inv_expansion + scalarinv_deriv(base_expansion, order=[o-1])
        base_expansion = base_expansion + scalarprod_deriv(norm_inv_expansion, a_expansion, order=[o-1])
        #     norm_expansion = norm_expansion + tensordot_deriv(vec_expansion, unit_expansion, axes=[-1, -1], order=[o],
        #                                                       shared=vec_expansion[0].ndim-1)
        #     norm_inv_expansion = norm_inv_expansion + scalarinv_deriv(norm_expansion, order=[o])
        # if len(unit_expansion) <= o:
        #     deriv_term = scalarprod_deriv(norm_inv_expansion, vec_expansion, order=[o])
        #     unit_expansion = unit_expansion + deriv_term

    # reexpand in terms of original vectors
    # raise Exception(vec_expansion)
    norm_expansion = [base_expansion[0]] + tensor_reexpand(
        vec_expansion[1:],
        base_expansion[1:],
        order + 1
        # shared=vec_expansion[0].ndim - 1
    )

    unit_expansion = [base_expansion[1]] + tensor_reexpand(
        vec_expansion[1:],
        base_expansion[2:],
        order
        # shared=vec_expansion[0].ndim - 1
    )

    return norm_expansion, unit_expansion

def vec_anglecos_deriv(A_expansion, B_expansion, order, unitized=False):
    if not unitized:
        A_expansion = vec_norm_unit_deriv(A_expansion, order)[1]
        B_expansion = vec_norm_unit_deriv(B_expansion, order)[1]
    shared = A_expansion[0].ndim - 1
    return tensordot_deriv(A_expansion, B_expansion, order=order, axes=[-1, -1], shared=shared)

def vec_cross_deriv(A_expansion, B_expansion, order):
    shared = A_expansion[0].ndim - 1
    base_shape = A_expansion[0].shape[:-1]
    e3 = np.broadcast_to(
        np.expand_dims(levi_cevita3, list(range(shared))),
        base_shape + (3, 3, 3)
    )
    return tensorops_deriv(
        B_expansion,
            [-1, shared],
        [e3],
            [-1, -1],
        A_expansion,
        order=order,
        shared=shared
    )

def vec_parallel_cross_norm_deriv(A_expansion, B_expansion, order):
    i3 = np.broadcast_to(np.eye(3)[np.newaxis], (3, 3, 3)).copy()
    i3[0, 0, 0] = 0
    i3[1, 1, 1] = 0
    i3[2, 2, 2] = 0

    base_shape = A_expansion[0].shape[:-1]
    shared = len(base_shape)
    i3 = np.broadcast_to(np.expand_dims(i3, list(range(shared))), base_shape + (3, 3, 3))

    pseudo_norm_2 = np.sqrt(tensorops_deriv(
        B_expansion,
            [-1, -1],
        [i3],
            [-1, -1],
        A_expansion,
        order=0,
        shared=shared
    )[0])

    if is_numeric(order):
        order = np.arange(order+1)

    A_exp = [
        vec_ops.vec_tensordot(A_expansion[o], pseudo_norm_2, axes=[-1, -1], shared=shared)
            if len(A_expansion) > o and not is_zero(A_expansion[o])  else
        0
        for o in order
    ]
    B_exp = [
        vec_ops.vec_tensordot(B_expansion[o], pseudo_norm_2, axes=[-1, -1], shared=shared)
            if len(B_expansion) > o and not is_zero(B_expansion[o]) else
        0
        for o in order
    ]
    pseudonorm_expansion = [a-b for a,b in zip(A_exp, B_exp)]
    return pseudonorm_expansion

def vec_anglesin_deriv(A_expansion, B_expansion, order, unitized=False, return_unit_vectors=True, planar=None, planar_threshold=1e-14):
    if is_numeric(order):
        max_order = order
    else:
        max_order = max(order)
    if not unitized:
        A_expansion = vec_norm_unit_deriv(A_expansion, max_order)[1]
        B_expansion = vec_norm_unit_deriv(B_expansion, max_order)[1]

    if not planar:
        vec_cross = vec_cross_deriv(A_expansion, B_expansion, order=max_order)
        if planar is None:
            planar = np.linalg.norm(vec_cross[0], axis=-1) <= planar_threshold
    else:
        vec_cross = None

    if np.all(planar):
        expansion = vec_parallel_cross_norm_deriv(A_expansion, B_expansion, order)
        if return_unit_vectors:
            return expansion, A_expansion
        else:
            return expansion
    elif np.any(planar):
        # gotta compute planar and nonplanar separately then remerge the tensors

        base_shape = A_expansion[0].shape[:-1]
        planar_pos = np.where(planar)
        nonplanar_pos = np.where(np.logical_not(planar))
        planar_A = [a[planar_pos] for a in A_expansion]
        planar_B = [b[planar_pos] for b in B_expansion]
        expansion = vec_parallel_cross_norm_deriv(planar_A, planar_B, order)
        planar_units = planar_A
        nonplanar_geoms = [v[nonplanar_pos] for v in vec_cross]
        norms, units = vec_norm_unit_deriv(nonplanar_geoms, order)

        final_tensors = [
            np.zeros(base_shape + e.shape[1:])
            for e in expansion
        ]
        for e,n,f in zip(expansion, norms, final_tensors):
            f[planar_pos] = e[planar_pos]
            f[nonplanar_pos] = n[nonplanar_pos]
        if return_unit_vectors:
            final_units = [
                np.zeros(base_shape + u.shape[1:])
                for u in units
            ]
            for e, n, f in zip(planar_units, units, final_units):
                f[planar_pos] = e[planar_pos]
                f[nonplanar_pos] = n[nonplanar_pos]

            return final_tensors, final_units
        else:
            return final_tensors
    else:
        norms, units = vec_norm_unit_deriv(vec_cross, order)
        if return_unit_vectors:
            return norms, units
        else:
            return norms

def vec_angle_deriv(A_expansion, B_expansion, order, unitized=False):
    if not unitized:
        units_A, A_expansion = vec_norm_unit_deriv(A_expansion, order)
        units_B, B_expansion = vec_norm_unit_deriv(B_expansion, order)
    cos_expansion = vec_anglecos_deriv(A_expansion, B_expansion, order, unitized=True)
    sin_expansion = vec_anglesin_deriv(A_expansion, B_expansion, order, unitized=True, return_unit_vectors=False)
    cos_sin_expansion = [
        np.moveaxis(np.array([c, s]), 0, -1)
        for c, s in zip(cos_expansion, sin_expansion)
    ]
    arctan_expansion = [
        np.array([
            -sin_expansion[0],
            cos_expansion[0]
        ]),
        np.array([
            [np.zeros_like(sin_expansion[0]), -np.ones_like(sin_expansion[0])],
            [np.ones_like(sin_expansion[0]),   np.zeros_like(cos_expansion[0])]
        ])
    ]
    # raise Exception(arctan_expansion)
    new_expansions = []
    for i, e in enumerate(arctan_expansion):
        for _ in range(i + 1):
            e = np.moveaxis(e, 0, -1)
        new_expansions.append(e)
    arctan_expansion = new_expansions
    angle_derivs = [np.arctan2(sin_expansion[0], cos_expansion[0])] + tensor_reexpand(
        cos_sin_expansion[1:],
        arctan_expansion,
        # axes=[-1, -1],
        order=order,
        # shared=arctan_expansion[0].ndim - 1
    )
    return angle_derivs

def vec_dihed_deriv(A_expansion, B_expansion, C_expansion, order, planar=None, planar_threshold=1e-14):
    # quick check

    if is_numeric(order):
        max_order = order
    else:
        max_order = max(order)

    axb_expansion = vec_cross_deriv(A_expansion, B_expansion, max_order)
    bxc_expansion = vec_cross_deriv(B_expansion, C_expansion, max_order)

    return vec_angle_deriv(axb_expansion, bxc_expansion, order, unitized=False)

def vec_plane_angle_deriv(A_expansion, B_expansion, C_expansion, D_expansion, order, planar=None, planar_threshold=1e-14):
    # quick check

    if is_numeric(order):
        max_order = order
    else:
        max_order = max(order)

    axb_expansion = vec_cross_deriv(A_expansion, B_expansion, max_order)
    cxd_expansion = vec_cross_deriv(C_expansion, D_expansion, max_order)

    return vec_angle_deriv(axb_expansion, cxd_expansion, order, unitized=False)


def nca_partition_terms(partition):
    """
    Computes the number of permutations for the non-commutative operation
    corresponding to the given partition

    :param cls:
    :param partition:
    :return:
    """
    _, counts = np.unique(partition, return_counts=True)

    # compute the reduced multinomial coefficient in stable (enough) form
    multinomial_num = np.flip(np.arange(1, 1 + np.sum(partition)))
    multinomial_denum = np.concatenate([
        np.arange(1, 1 + j)
        for j in partition
    ])
    multi_redux_terms = np.concatenate([
        np.arange(1, 1 + j)
        for j in counts
    ])
    full_denom = np.flip(np.sort(np.concatenate([multinomial_denum, multi_redux_terms])))
    numl = len(multinomial_num)
    denl = len(full_denom)
    if numl > denl:
        full_denom = np.concatenate([full_denom, np.ones(numl - denl, dtype=full_denom.dtype)])
    elif denl > numl:
        multinomial_num = np.concatenate([multinomial_num, np.ones(denl - numl, dtype=multinomial_num.dtype)])

    return np.prod(multinomial_num / full_denom)

def nca_symmetrize(tensor, partition,
                   shared=None,
                   identical=True,
                   reweight=None):
    from ..Combinatorics import UniquePermutations
    perm_counter = len(partition)
    perm_idx = []  # to establish the set of necessary permutations to make things symmetric
    for i in partition:
        if (not identical) or i > 1:
            perm_idx.extend([perm_counter] * i)
            perm_counter -= 1
        else:
            perm_idx.append(perm_counter)

    # sometimes we overcount, so we factor that out here
    perm_inds, _ = UniquePermutations(perm_idx).permutations(return_indices=True)
    if reweight or (reweight is None and identical):
        nterms = nca_partition_terms(partition)
        overcount = len(perm_inds) / nterms
        tensor = tensor / overcount

    if shared is None:
        perm_inds = [
            list(p) + list(range(len(p), tensor.ndim))
            for p in perm_inds
        ]
    else:
        l = list(range(shared))
        perm_inds = [
            l + [shared + pp for pp in p] + list(range(shared+len(p), tensor.ndim))
            for p in perm_inds
        ]

    return sum(
        tensor.transpose(p)
        for p in perm_inds
    )

def nca_partition_dot(partition, A_expansion, B_expansion, axes=None, shared=None, symmetrize=True):
    if axes is None:
        axes = [-1, (0 if shared is None else shared)]
    if len(B_expansion) <= len(partition) - 1:
        return 0
    B = B_expansion[len(partition) - 1]
    if is_numeric(B) and B == 0:
        return 0
    a_ax, b_ax = [[x] if is_numeric(x) else x for x in axes]
    b_ax = [B.ndim + b if b < 0 else b for b in b_ax]
    for i in reversed(partition):
        if len(A_expansion) <= i - 1:
            return 0
        A = A_expansion[i - 1]
        if is_numeric(A) and A == 0:
            return 0
        if shared is None:
            B = np.tensordot(A, B, axes=[a_ax, b_ax])
        else:
            B = vec_ops.vec_tensordot(A, B, axes=[a_ax, b_ax], shared=shared)
        b_ax = [min(B.ndim - 1, b + A.ndim - 1 - shared) for b in b_ax]

    if symmetrize:
        B = nca_symmetrize(B, partition, shared=shared)
    return B

def nca_partition_prod(partition, A_expansion, shared=None, symmetrize=True):
    i = partition[0]
    if len(A_expansion) <= i - 1:
        return 0
    A = A_expansion[i - 1]
    if is_numeric(A) and A == 0:
        return 0

    B = A
    for i in partition[1:]:
        if len(A_expansion) <= i - 1:
            return 0
        A = A_expansion[i - 1]
        if is_numeric(A) and A == 0:
            return 0
        if shared is not None:
            axes = [list(range(shared, A.ndim)), list(range(shared, B.ndim))]
        else:
            axes = [list(range(A.ndim)), list(range(B.ndim))]
        B = vec_ops.vec_outer(A, B, axes=axes)

    if symmetrize:
        B = nca_symmetrize(B, partition, shared=shared)
    return B

def tensor_reexpand(derivs, vals, order=None, axes=None):
    from ..Combinatorics import IntegerPartitioner
    terms = []
    if order is None:
        order = len(vals)

    shared = derivs[0].ndim - 2
    if is_numeric(order): order = list(range(1, order+1))

    for o in order:
        term = sum(
            nca_partition_dot(p, derivs, vals, axes=axes, shared=shared)
            for parts in IntegerPartitioner.partitions(o)
            for p in parts
        )
        terms.append(term)

    return terms

def optimizing_transformation(expansion, order):
    V = expansion

    zero_order = 0
    for v in V:
        if is_zero(v):
            zero_order += 1
        else:
            break

    if zero_order > 2:
        raise ValueError("can't optimize without gradient or Hessian (tensor inverses not implemented)")

    W = V[zero_order]
    if zero_order == 0:
        w = W
        Q = [np.eye(V[1].shape[0])]
    elif zero_order == 1 and np.allclose(np.diag(np.diag(W)), W):
        Q = [np.eye(V[1].shape[0])]
        w = np.diag(W)
    else:
        Q = [np.linalg.inv(W)]
        w = np.ones(len(W))

    for o in range(1, order + 1):
        V_rem = tensor_reexpand(Q, V, [o], axes=[-1, 0])[-1]
        w = w[np.newaxis]
        new_Q = -V_rem / ((o + 2) * w)
        Q = Q + [new_Q]

    return Q

def apply_nca_multi_ops(partition, expansions, ops, shared):
    terms = [
        e[p] if len(e) > p else 0
        for e,p in zip(expansions, partition)
    ]
    if any(is_numeric(t) and t == 0 for t in terms): return 0

    if shared is None: shared = 0

    A = terms[0]
    # d = partition[0]
    d = A.ndim - 2 - shared
    scaling = 2 if len(partition) == 3 and tuple(partition) == (1, 0, 1) else 1
    for B,p,(op, axes, contract) in zip(terms[1:], partition[1:], ops):
        a_ax, b_ax = [[x] if is_numeric(x) else x for x in axes]
        a_ax = [a if a >= 0 else (A.ndim + a - d) for a in a_ax]
        b_ax = [b if b >= 0 else (B.ndim + b - p) for b in b_ax]
        if contract: # axes disappear, so we just account for the shifts
            axes = [[x+d for x in a_ax], [x+p for x in b_ax]]
            deriv_axis = A.ndim - len(a_ax)
        else: # axes appeared, so we need to include those in the product
              # we _forced_ axes to be at the end of the arrays since it was too hard
              # to keep track of them otherwise...
              # actually I guess I could have put the derivative axes at the end...
              # and then the axes would never change...but that has other complications
            axes = [
                [shared + i for i in range(d)] + [x+d for x in a_ax],
                [shared + i for i in range(p)] + [x+p for x in b_ax]
            ]
            deriv_axis = A.ndim
        if shared == 0:
            A = op(A, B, axes=axes)
        else:
            A = op(A, B, axes=axes, shared=shared)
        # next we need to move all of the derivative axes in the second tensor to the beginning
        # so we can symmetrize

        for i in range(p):
            A = np.moveaxis(A, deriv_axis+i, shared)
        d += B.ndim - 2 - shared
    partition = [p for p in partition if p > 0]
    if len(partition) > 1:
        A = nca_symmetrize(A, partition, shared=shared, identical=False)
    return A
def nca_multi_op_order_deriv(partition_generator, order, expansions, ops, shared):
    full = None
    for part in partition_generator.get_terms(order):
        term = apply_nca_multi_ops(part, expansions, ops, shared)
        if full is None:
            full = term
        else:
            full = full + term
    return full

def nca_canonicalize_multiops(expansion_op_pairs):
    if len(expansion_op_pairs) % 2 == 0: raise ValueError("invalid number of operations/expansions")

    expansions = []
    ops = []
    for i, t in enumerate(expansion_op_pairs):
        if i % 2 == 0:
            expansions.append([np.asanyarray(A) for A in t])
        else:
            op = None
            axes = None
            contract = None
            if isinstance(t, dict):
                op = t.get('op', None)
                axes = t.get('axes', None)
                contract = t.get('contract', None)
            elif isinstance(t, str):
                op = t
            elif is_numeric(t[0]) or is_numeric(t[0][0]):
                axes = t
            else:
                op = t[0]
                if len(t) > 1: axes = t[1]
                if len(t) > 2: contract = t[2]

            if op is None and axes is not None:
                op = '.'

            if isinstance(op, str):
                if op in {'.', 'dot', 'contract', 'inner'}:
                    op = _contract
                    if axes is None: axes = [-1, 0]
                    if contract is None: contract = True
                elif op in {'x', 'prod', 'product', 'outer'}:
                    op = _product
                    # if axes is None: axes = None
                    if contract is None: contract = False
                else:
                    raise ValueError(f"can't canonicalize operation {op}")

            if contract is None:
                contract = op is _contract

            ops.append([op, axes, contract])

    return expansions, ops

def _contract(a, b, axes=None, shared=None):
    if shared is not None:
        res = vec_ops.vec_tensordot(a, b, axes=axes, shared=shared)
    else:
        res = np.tensordot(a, b, axes=axes)
    return res
def _product(a, b, axes=None, shared=None):
    return vec_ops.vec_outer(a, b, axes=axes)

def tensorops_deriv(
        *expansion_op_pairs,
        order,
        shared=None
):
    from ..Combinatorics import SymmetricGroupGenerator
    expansions, ops = nca_canonicalize_multiops(expansion_op_pairs)

    if isinstance(order, int):
        order = list(range(order+1))

    partition_generator = SymmetricGroupGenerator(len(expansions))
    derivs = [
        nca_multi_op_order_deriv(partition_generator, o, expansions, ops, shared)
        for o in order
    ]

    return derivs