"""
File Purpose: Component, ComponentList, ComponentDimension, ComponentHaver
"""

from .dimension_tools import (
    DimensionValue, DimensionValueList, string_int_lookup,
    Dimension, DimensionHaver,
)
from ..errors import ComponentKeyError, ComponentValueError, InputError
from ..tools import is_integer, BasisXYZ
from ..tools.xarray_tools import xarray_vectors  # so we can update _BASIS there.

### --------------------- Component & ComponentList --------------------- ###

class Component(DimensionValue):
    '''a single vector component coordinate (e.g. x, y, or z).
    Knows how to be converted to str and int. e.g. 'y' or 1.

    s: string or None
        str(self) --> s. if None, cannot convert to str.
    i: int or None
        int(self) --> i. if None, cannot convert to int.

    Note: not related to "coordinates", e.g. location in space.
    '''
    _dimension_key_error = ComponentKeyError
    _dimension_value_error = ComponentValueError

    def __init__(self, s=None, i=None):
        super().__init__(s, i)


class ComponentList(DimensionValueList):
    '''list of vector component coordinates.'''
    _dimension_key_error = ComponentKeyError
    _dimension_value_error = ComponentValueError
    value_type = Component

    @classmethod
    def from_strings(cls, strings):
        '''return ComponentList from iterable of strings. (i will be determined automatically.)'''
        return cls(Component(s, i) for i, s in enumerate(strings))

    @classmethod
    def from_array(cls, array):
        '''return ComponentList from 0D or 1D array.
        array values should be Component objects, strings, or ints.
            Component objects --> use as-is
            strings --> if 'x', 'y', or 'z', use XYZ.get(string).
                        Else, use Component(s, i=i) for i in range(len(array))
            ints --> if 0, 1, or 2, use XYZ.get(int).
                     Else, use Component(i=i) for i in array
        '''
        # [TODO] use super() instead of repeating code from DimensionValueList.from_array?
        if array.ndim == 0:
            values = [array.item()]
        elif array.ndim == 1:
            values = array.values
        else:
            errmsg = f'{cls.__name__}.from_array expects array.ndim=0 or 1; got ndim={array.ndim}'
            raise DimensionalityError(errmsg)
        result = []
        for i, v in enumerate(values):
            if isinstance(v, cls.value_type):
                result.append(v)
            elif (isinstance(v, str) or isinstance(v, int)) and v in cls.value_type.XYZ:
                result.append(cls.value_type.XYZ.get(v))
            elif isinstance(v, str):
                result.append(cls.value_type(v, i=i))
            elif is_integer(v):
                result.append(cls.value_type(i=v))
            else:
                errmsg = (f'{cls.__name__}.from_array got unexpected value type at index {i}.\n'
                            f'Expected {cls.value_type.__name__}, str, or int; got value={v!r}.')
                raise InputError(errmsg)
        return cls(result)


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

    XYZ: ComponentList with length 3
        contains [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: string_int_lookup({X:(Y,Z), Y:(Z,X), Z:(X,Y)}))
    '''
    def __init__(self, XYZ):
        if not isinstance(XYZ, ComponentList) or len(XYZ) != 3:
            errmsg = f'ComponentBasisXYZ requires XYZ to be ComponentList of length 3; got {XYZ!r}'
            raise InputError(errmsg)
        super().__init__(XYZ)
        self.YZ_FROM_X = string_int_lookup(self.YZ_FROM_X)
    
    def get(self, x):
        '''return proper component label (i.e. from self.XYZ) given x.
        x can be any valid component specifier, e.g. int, str, Component.
        Equivalent to self.XYZ.get(x).
        '''
        return self.XYZ.get(x)


XYZ = ComponentList.from_strings('xyz')
BASIS_XYZ = ComponentBasisXYZ(XYZ)
_BASIS = BASIS_XYZ

X, Y, Z = _BASIS.XYZ
YZ_FROM_X = _BASIS.YZ_FROM_X

Component.YZ_FROM_X = YZ_FROM_X  # add to Component class for convenience.
Component.XYZ = XYZ  # add to Component class for convenience.

XHAT = _BASIS.XHAT
YHAT = _BASIS.YHAT
ZHAT = _BASIS.ZHAT

# UPDATE tools._BASIS:
xarray_vectors._BASIS = _BASIS


### --------------------- ComponentDimension, ComponentHaver --------------------- ###

class ComponentDimension(Dimension, name='component', plural='components',
                     value_error_type=ComponentValueError, key_error_type=ComponentKeyError):
    '''component dimension, representing current value AND list of all possible values.
    Also has various helpful methods for working with this Dimension.
    '''
    pass  # behavior inherited from Dimension.


@ComponentDimension.setup_haver
class ComponentHaver(DimensionHaver, dimension='component', dim_plural='components'):
    '''class which "has" a ComponentDimension. (ComponentDimension instance will be at self.component_dim)
    self.component stores the current vector component (possibly multiple). If None, use self.components instead.
    self.components stores "all possible vector components" for the ComponentHaver.
    Additionally, has various helpful methods for working with the ComponentDimension,
        e.g. current_n_component, iter_components, take_component.
        See ComponentDimension.setup_haver for details.

    components defaults to XYZ (==ComponentList.from_strings('xyz'))
    '''
    def __init__(self, *, component=None, components=XYZ, **kw):
        super().__init__(**kw)
        if components is not None: self.components = components
        self.component = component
