"""
File Purpose: tools related to xarray vector arithmetic along 'component' dimension,
probably with 'x', 'y', 'z' components.

NOTE: functions with multiple vector inputs here are not registered here with pcAccessor.
    See vector_arithmetic.py for similar versions which ARE registered with pcAccessor.
    (Those versions might eventually be updated to reference-frame-aware checks,
    taking Component machinery under consideration.
    The versions here will instead never do any implicit reference frame rotations.)
"""
import numpy as np
import xarray as xr

from .xarray_accessors import pcAccessor
from .xarray_dimensions import (
    take_along_dimension, join_along_dimension,
    xarray_assign,
)
from .xarray_indexing import xarray_sel
from ..sentinels import UNSET
from ...errors import InputError

### --------------------- Allow vector tools without Component --------------------- ###
# Tools should all work properly whether _BASIS is still strings, or proper Component objects.
# It will be strings if ported these tools via copy-pasting.
# It will be Component objects if used inside PlasmaCalcs, as components.py updates _BASIS.

class BasisXYZ():
    '''simple class to hold X, Y, Z component details.
    (Assumed to be immutable after creation)

    XYZ: list-like with length 3
        iterable containing [X, Y, Z] component labels.

    Instances of this class provide attributes:
        X: x component label
        Y: y component label
        Z: z component label
        XHAT: x unit vector, as a DataArray of [1,0,0] with 'component' values XYZ
        YHAT: y unit vector, as a DataArray of [0,1,0] with 'component' values XYZ
        ZHAT: z unit vector, as a DataArray of [0,0,1] with 'component' values XYZ
        YZ_FROM_X: dict telling (y,z) given x, for right-handed coord system x,y,z
            (internally here: {X:(Y,Z), Y:(Z,X), Z:(X,Y)})
    '''
    def __init__(self, XYZ):
        self.XYZ = XYZ
        self.X = self.XYZ[0]
        self.Y = self.XYZ[1]
        self.Z = self.XYZ[2]
        self.XHAT = xr.DataArray([1,0,0], coords={'component': self.XYZ})
        self.YHAT = xr.DataArray([0,1,0], coords={'component': self.XYZ})
        self.ZHAT = xr.DataArray([0,0,1], coords={'component': self.XYZ})
        self.YZ_FROM_X = {
            self.X: (self.Y, self.Z),
            self.Y: (self.Z, self.X),
            self.Z: (self.X, self.Y),
        }

    def get(self, x):
        '''return proper component label (i.e. from self.XYZ) given x.
        This is the first match from XYZ which equals x. Crash if no match.
        '''
        for xc in self.XYZ:
            if xc == x:
                return xc

    def get_hat(self, x):
        '''return self.XHAT, self.YHAT, or self.ZHAT depending on str(self.get(x)).'''
        _x_input = x
        x = str(self.get(x))
        if x == str(self.X):
            return self.XHAT
        elif x == str(self.Y):
            return self.YHAT
        elif x == str(self.Z):
            return self.ZHAT
        else:
            errmsg = (f'{type(self).__name__}.hat(x={_x_input!r}), but str(self.get(x))=={x!r} not found in: '
                      f'str(self.X)={str(self.X)!r}, str(self.Y)={str(self.Y)!r}, str(self.Z)={str(self.Z)!r}.')
            raise InputError(errmsg)

    def __repr__(self):
        return f'{type(self).__name__}(XYZ={self.XYZ})'

_BASIS = BasisXYZ(['x', 'y', 'z'])  # default basis, using strings
# CAUTION: don't rely on _BASIS for kwarg defaults, only use it inside functions.

def _get_basis(_basis=UNSET):
    '''return proper BasisXYZ to use.
    _basis: UNSET or BasisXYZ
        tells basis to use.
        UNSET --> default to xarray_vectors._BASIS
    '''
    return _BASIS if _basis is UNSET else _basis


### --------------------- Independent of _BASIS --------------------- ###

# @pcAccessor.register('dot_product', aliases=['dot'])  # registered in vector_arithmetic.py instead
def xarray_dot_product(A, B):
    '''return dot product of vectors A and B, assuming vector components along the dimension 'component'.'''
    # note: need to ensure to only sum across matching components
    #    (can't just take_along_dimension('component') without checking if components match)
    # also note: (A * B).sum('component') fails when component is a non-dimension coordinate,
    #    e.g. when there is only 1 component.
    # [TODO] probably a more concise way to write this function, using xarray_promote_dim.
    A_dim = 'component' in A.dims
    B_dim = 'component' in B.dims
    # case: both have 'component' dimension. Easy - let xarray handle all the bookkeeping.
    if A_dim and B_dim:
        result = A * B
        return result.sum('component', skipna=False)  # skipna=False --> result NaN wherever any summand is NaN
    # case: exactly one has 'component' dimension but the other does not.
    elif A_dim or B_dim:
        if A_dim:
            a, b = A, B
        else:
            a, b = B, A
        # now, 'a' has 'component' dimension while 'b' does not.
        if b.coords['component'].size != 1:
            raise NotImplementedError('dot_product, when component is a non-dimension coordinate with size > 1')
        idx = np.nonzero(a.coords['component'].values == b.coords['component'].values)[0]
        if len(idx) == 0:  # no matching components. Result is 0.
            return xr.zeros_like(b).drop_vars('component')  # drop_vars --> result has no 'component' coord.
        else:
            assert len(idx) == 1
            i = idx[0]
            result = a.isel(component=i) * b
            return result.drop_vars('component')  # drop_vars --> result has no 'component' coord.
    # case: A and B both do not have 'component' dimension.
    else:
        # first ensure A and B have the same 'component' coordinate, then it's easy, just multiply them.
        if not (A.coords['component'].size == 1 == B.coords['component'].size):
            raise NotImplementedError('dot_product, when component is a non-dimension coordinate with size > 1')
        A_comps = A.coords['component'].values
        B_comps = B.coords['component'].values
        if np.all(A_comps == B_comps):
            result = A * B
            return result.drop_vars('component')  # drop_vars --> result has no 'component' coord.
        else:
            return xr.zeros_like(A).drop_vars('component')  # drop_vars --> result has no 'component' coord.
    assert False, "if reached this line, made coding mistake."

@pcAccessor.register('magnitude', aliases=['mag', 'mod'])
def xarray_magnitude(A, *, squared=False):
    '''return vector magnitude of A, assuming vector components along the dimension 'component'.
    squared: bool, default False
        if True, return |A|**2 instead of |A|.
        [EFF] to get |A|**2, when |A| is not needed,
            magnitude(A, squared=True) is more efficient than magnitude(A)**2
    '''
    result = xarray_dot_product(A, A)
    return result if squared else result**0.5

@pcAccessor.register('unit_vector', aliases=['hat', 'dir', 'direction'])
def xarray_unit_vector(A):
    '''return a unit vector in the direction of A. Equivalent: A / |A|'''
    return A / xarray_magnitude(A)

@pcAccessor.register('rmscomps')
def xarray_rmscomps(A):
    '''return root mean squared of components of A.
    E.g., rmscomps(A) --> sqrt((Ax^2 + Ay^2 + Az^2) / 3), if A has 3 components.
    '''
    num_comps = A.coords['component'].size  # number of components
    mag2 = xarray_magnitude(A, squared=True)
    return np.sqrt(mag2 / num_comps)

# intentionally no @pcAccessor; see xarray_perp below instead.
def xarray_take_perp_to(B, A):
    '''return the component of A perpendicular to B. Equivalent: A - (A dot Bhat) Bhat.
    Note that B is the first argument.
    '''
    Bhat = B / xarray_magnitude(B)
    return A - xarray_dot_product(A, Bhat) * Bhat

# intentionally no @pcAccessor; see xarray_parallel below instead.
def xarray_take_parallel_to(B, A):
    '''return the component of A parallel to B. Equivalent: (A dot Bhat) Bhat.
    Note that B is the first argument.
    '''
    Bhat = B / xarray_magnitude(B)
    return xarray_dot_product(A, Bhat) * Bhat

# @pcAccessor.register('perp')  # in vector_arithmetic.py instead
def xarray_perp(A, B):
    '''return the component of A perpendicular to B. Equivalent: A - (A dot Bhat) Bhat.'''
    return xarray_take_perp_to(B, A)

# @pcAccessor.register('parallel', aliases=['par'])  # in vector_arithmetic.py instead
def xarray_parallel(A, B):
    '''return the component of A parallel to B. Equivalent: (A dot Bhat) Bhat.'''
    return xarray_take_parallel_to(B, A)

# @pcAccessor.register('perpmod', aliases=['perpmag'])  # in vector_arithmetic.py instead
def xarray_perpmod(A, B):
    '''return the magnitude of the component of A perpendicular to B. Equivalent: |A - (A dot Bhat) Bhat|'''
    return xarray_magnitude(perp(A, B))

# @pcAccessor.register('parmod', aliases=['parmag'])  # in vector_arithmetic.py instead
def xarray_parmod(A, B, *, signed=False):
    '''return the magnitude of the component of A parallel to B. Equivalent: |(A dot Bhat) Bhat|
    signed: bool; whether to return signed result instead, i.e. (A dot Bhat), without taking abs.
    '''
    result = xarray_dot_product(A, xarray_unit_vector(B))
    return result if signed else np.abs(result)



### --------------------- Related to / Depends on _BASIS --------------------- ###


@pcAccessor.register('take_xyz')
def xarray_take_xyz(A, *, drop=False, as_dict=False):  
    '''return [Ax, Ay, Az] from selecting 'x', 'y', 'z' along 'component' dimension.
    Crashes if any components are missing.

    drop: bool
        whether to drop x, y, z 'component' coord from results.
    as_dict: bool
        if True, return dict(x=Ax, y=Ay, z=Az) instead of list [Ax, Ay, Az].
    '''
    x = xarray_sel(A, component='x', drop=drop)
    y = xarray_sel(A, component='y', drop=drop)
    z = xarray_sel(A, component='z', drop=drop)
    return dict(x=x, y=y, z=z) if as_dict else [x, y, z]

@pcAccessor.register('join_xyz', aliases=['join_yz'])
def xarray_join_xyz(x, y, z, *, _basis=UNSET):
    '''return vector with components x, y, z along 'component' dimension.
    (Call as xarray_join_xyz(x, y, z), or x.pc.join_yz(y,z), or x.pc.join_xyz(y,z))
    Result has 'component' coord with values 'x', 'y', 'z'
        (or Component objects representing 'x', 'y', 'z', depending on _basis.)

    x, y, z: xarray.DataArray or Dataset objects
        vector components to use.
    _basis: UNSET or BasisXYZ
        tells basis from which to infer X, Y, Z component labels.
        UNSET --> default to xarray_vectors._BASIS
    '''
    if _basis is UNSET: _basis = _BASIS
    x = x.assign_coords(component=_basis.X)
    y = y.assign_coords(component=_basis.Y)
    z = z.assign_coords(component=_basis.Z)
    return join_along_dimension('component', [x, y, z])

# @pcAccessor.register('angle_xy', aliases=...)  # in vector_arithmetic.py instead
def xarray_angle_xy(A):   # okay this one doesn't actually depend on _BASIS but it comes pretty close...
    '''return angle between +xhat and A, in the xy plane, in radians.
    A should be a DataArray (or Dataset) with x and y in 'component' dimension.
    A can be any vector (does not need to be a unit vector, but it can be.)
    '''
    x = xarray_sel(A, component='x', drop=True)
    y = xarray_sel(A, component='y', drop=True)
    return np.arctan2(y, x)

# @pcAccessor.register('angle_xy_to_hat', ...)  # in vector_arithmetic.py instead
def xarray_angle_xy_to_hat(A, *, _basis=UNSET):
    '''return unit vector u, given angle [radians] between +xhat and u in the xy plane.
    The result has 3 components (x,y,z) with value equal to 0 for z component.
    Equivalent: cos(A) * xhat + sin(A) * yhat.

    _basis: UNSET or BasisXYZ
        tells basis from which to infer XHAT and YHAT
        UNSET --> default to xarray_vectors._BASIS
    '''
    if _basis is UNSET: _basis = _BASIS
    return np.cos(A) * _basis.XHAT + np.sin(A) * _basis.YHAT

# @pcAccessor.register('cross_component')  # in vector_arithmetic.py instead
def xarray_cross_component(A, B, x, *, yz=None, missing_ok=False, _basis=UNSET):
    '''return x component of A cross B, given A and B which have values for y and z 'component'.

    x: str (or any other valid component specifier)
        probably 'x', 'y', or 'z'; tells which component of A cross B to return.
        will be "normalized" via _basis.get(x).
    A, B: xarray.DataArray
        vectors to take cross product of.
        must include 'component' dimension including coordinates y and z.
    yz: None or iterable of two strings objects
        the other two components; (x,y,z) should form a right-handed coordinate system.
        if not provided, infer from x.
    missing_ok: bool, default False
        whether it is okay for 'component' dimension to be missing y or z components, of A or B.
        if True, treat any missing components as 0.
    _basis: UNSET or BasisXYZ
        tells basis from which to normalize x and to infer YZ_FROM_X mapping.
        UNSET --> default to xarray_vectors._BASIS
    '''
    if _basis is UNSET: _basis = _BASIS
    x = _basis.get(x)
    if yz is None: yz = _basis.YZ_FROM_X[x]
    default = 0 if missing_ok else UNSET
    Ay, Az = take_along_dimension('component', A, yz, default=default)
    By, Bz = take_along_dimension('component', B, yz, default=default)
    return xarray_assign(Ay * Bz - Az * By, coords={'component': x})

# @pcAccessor.register('cross_product')  # in vector_arithmetic.py instead
def xarray_cross_product(A, B, *, components=None, _basis=UNSET):
    '''return cross product of vectors A and B, along dimension 'component'.
    If A or B missing any components, treat them as 0.

    components: None or iterable of str (or any other valid component specifier)
        tells which components to get.
        None --> get all components (XYZ)
        will be "normalized" via _basis.get(x) for each component (see xarray_cross_component)
    _basis: UNSET or BasisXYZ
        tells basis from which to infer XYZ component labels.
        UNSET --> default to xarray_vectors._BASIS
    '''
    if _basis is UNSET: _basis = _BASIS
    if components is None:
        components = _basis.XYZ
    result = [xarray_cross_component(A, B, x, missing_ok=True) for x in components]
    return join_along_dimension('component', result)
