# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06_axes.ipynb.

# %% auto 0
__all__ = ['npmixin', 'naxes']

# %% ../nbs/06_axes.ipynb 4
from typing import (
    TypeAlias, Optional, Iterable, Union, Tuple,
)
from atyp import IntQ, BoolQ, FloatQ

# %% ../nbs/06_axes.ipynb 5
from indoc import indoc
from fuzstr import fuzstr
from nymstr import nymstr
from synset import synset

from aspec import aspec

# %% ../nbs/06_axes.ipynb 6
from indoc import indoc
from fuzstr import fuzstr
from nymstr import nymstr
from synset import synset

from aspec import aspec

# %% ../nbs/06_axes.ipynb 7
from .atyp import AxisKey, IterAxes, IterAxesQ
from .copy import safecopy
from .axis import naxis
from .errs import AxisError

# %% ../nbs/06_axes.ipynb 9
class npmixin:
    @property
    def sizes(self) -> Optional[Tuple[int]]: 
        '''NumPy's `np.ndarray.shape` property if axes is subclassed (i.e. sized), else None'''
        shape = getattr(self, 'shape', None)
        return shape

    def asize(self, i) -> IntQ:
        '''Get `self.axes[i]` size if available (i.e. position / location in `self.axes`)'''
        return self.sizes[i] if self.sizes and i < len(self.sizes) else None

    def astr(self, i: int) -> str:
        '''Make `self.axes[i]` str (with size if available) representing `self.axes[i]`'''
        name = self.anames[i]
        size = self.asize(i)
        return name if size is None else f'{size} {name}' 

    @classmethod 
    def makeaxes(cls, ndim: int) -> IterAxes:
        return tuple(naxis(f'a{i}', axis=i) for i in range(ndim))

    def baseaxes(self) -> IterAxes:
        return safecopy(self.axes) if self.axes else self.makeaxes(self.ndim)
    
    @classmethod
    def safeidxs(cls, *axes: Iterable[str], **kwargs) -> IterAxesQ:
        axisrepr = kwargs.get('axisrepr', False)
        if axes and isinstance(axes[0], (list, tuple)) and not naxis.isinst(axes[0]):
            axes = axes[0]
        dims = list(axes or [])
        for i, axis in enumerate(dims):
            dims[i] = axis = naxis.fromstr(axis, axes=dims, **{**kwargs, **dict(fullrepr=axisrepr)})
        axes = (list(dims) or None)
        return axes   

# %% ../nbs/06_axes.ipynb 11
class naxes(npmixin, aspec):
    aspec: tuple = ('name', 'axes', )

    def __init__(
        self, *axes: IterAxes,
        icase:  BoolQ   = True,  # Case insensitive
        ispace: BoolQ   = True,  # Space insensitive
        iunder: BoolQ   = True,  # Underscore insensitive
        cutoff: FloatQ  = 0.8,   # Fuzzy cutoff        
        fullrepr: BoolQ = False, # Full representation
        axisrepr: BoolQ = False, # Whether to use axis fullreprs
        **kwargs
    ):
        cls = type(self)
        
        # Put all variables in a dictionary for easy access
        kwargs.update(
            icase=icase, ispace=ispace, iunder=iunder, cutoff=cutoff,
            fullrepr=fullrepr, axisrepr=axisrepr
        )

        # Make sure each axis has a unique index
        if axes and isinstance(axes[0], (list, tuple)) and not naxis.isinst(axes[0]):
            axes = axes[0]
        axes = cls.safeidxs(*axes, **kwargs)
        name = kwargs.pop('name', 'naxes')
        
        self._str = fuzstr(name, **kwargs)
        self._axs = axes
        self.fullrepr = fullrepr

    @property
    def name(self): return self._str
    @property
    def axes(self): return self._axs or []
    @property
    def naxes(self) -> int: return len((self.axes or []))
    @property
    def anames(self) -> Tuple[str]:
        '''Tuple of axis names if axes are named, else None'''
        return tuple(str(getattr(ax, 'name', ax)) for ax in self.axes) if self.axes else None
    
    def astr(self, i: int) -> str:
        '''Make `self.axes[i]` str (with size if available) representing `self.axes[i]`'''
        return self.anames[i]

    @property
    def astrs(self) -> str:
        '''Make axes str representing `self.axes`'''
        astrs = [self.astr(i) for i in range(self.naxes)]
        astrs = ', '.join(astrs)
        return f'({astrs})'
    
    def __str__(self): return f'{self.name}{self.astrs}'
    def __repr__(self): return f'{self.name}{self.axes}' if self.fullrepr else str(self)    
    def __iter__(self): return iter(self.axes)
    def __len__(self): return len(self.axes)
    def __contains__(self, axis: naxis): return axis in self.axes
    
    def __getaxis__(self, key: AxisKey) -> 'naxis':
        '''Get axis by key
        
        Notes
        -----
        This is the main method for getting an axis by key and is used by `__getitem__` and `getaxis`.
            It is provided as a dunder method so that it can be used as a method in `NamedArray`
            which subclasses `np.ndarray`.
        '''
        if isinstance(key, int):
            return self.axes[key]
        elif isinstance(key, naxis):
            return self.__getaxis__(key.name)
        elif isinstance(key, (str)):
            return next((ax for ax in self.axes if ax == key), None)
        elif isinstance(key, (synset, nymstr, fuzstr)):
            return next((ax for ax in self.axes if ax == key), None)
        else:
            raise AxisError(key, type(self))

    def getaxis(self, key: AxisKey) -> 'naxis':
        '''Get axis by key
        
        Notes
        -----
        This is an alias for `__getaxis__`. It is provided for convenience and so that it 
            can be used as a method in `NamedArray` which subclasses `np.ndarray`.    
        '''
        return self.__getaxis__(key)

    def aidx(self, key: AxisKey) -> naxis:
        '''Get axis index in `self.axes` by key'''
        ax = self.__getaxis__(key)
        return self.axes.index(ax)
    
    def index(self, key: AxisKey) -> 'naxis':
        '''Alias for `aidx`'''
        return self.aidx(key)

    def __getitem__(self, key: AxisKey) -> 'naxis':
        '''Get axis by key
        
        Notes
        -----
        This is an alias for `__getaxis__`. It is provided for convenience
            but it will be overwritten when subclassed with `np.ndarray`.
        '''
        return self.__getaxis__(key)

    def __taxes__(self, *order: AxisKey) -> IterAxes:
        '''Transpose axes by order
        
        Notes
        -----
        This is the main method for transposing axes by order and is used by `transpose`.
            It is provided as a dunder method so that it can be used as a method in `NamedArray`
            which subclasses `np.ndarray`.
        '''
        axes = safecopy(self.axes)
        # check input and convert to axes
        update_axes = [self.__getaxis__(key).copy() for key in order]
        # gather the axes that are not in the provided order
        needed_axes = [ax for ax in axes if ax not in update_axes]
        # the new order of axes is the updated axes followed by the needed axes
        new_order = update_axes + needed_axes        
        # rearrange axes according to the new order
        for i, ax in enumerate(new_order):
            ax.axis = i
        axes = new_order
        return axes
    
    def __setitem__(self, key: AxisKey, axis: naxis):
        try:
            prev = self.__getaxis__(key)
            self._ax[prev.axis] = axis
        except AxisError as err:
            if isinstance(key, int) and 0 <= key < self.naxes - 1: 
                self._ax[key] = axis
            else: raise err
        except KeyError as err:
            if isinstance(key, int) and 0 <= key < self.naxes - 1: 
                self._ax[key] = axis
            else: raise err
        except IndexError as err:
            if isinstance(key, int) and 0 <= key < self.naxes - 1: 
                self._ax[key] = axis
            else: raise err

    def append(self, axis: naxis):
        new = axis
        if isinstance(new, naxis):
            pass
        elif isinstance(new, str):
            new = naxis.fromstr(new, axes=self.axes)
        elif isinstance(new, int):
            new = naxis(f'a{new}', axis=new)
        else:
            raise ValueError(
                f'Expected new axis ({axis}) as int, str, '
                f'or naxis got {type(new)} instead.'
            )
        if axis not in self.axes:
            return self._ax.append(naxis)
        raise ValueError(f'New Axis ({axis}) already in axes')

    def transpose(self, *order: AxisKey) -> IterAxes:
        return self.__taxes__(*order)
    
    def __eq__(self, other):
        if isinstance(other, type(self)):
            return self.axes == other.axes
        elif isinstance(other, Iterable):
            return self.axes == other
        return False


