# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_narr.ipynb.

# %% auto 0
__all__ = ['narr']

# %% ../nbs/07_narr.ipynb 4
import numpy as np

from typing import (overload, Tuple,)
from atyp import ClsStrQ

# %% ../nbs/07_narr.ipynb 5
from .atyp import AxisKey, IterAxesQ
from .copy import safecopy
from .strs import first_line, first_digit_loc
from .axis import naxis
from .axes import naxes
from .errs import DimsError

# %% ../nbs/07_narr.ipynb 7
class narr(np.ndarray, naxes):
    name: ClsStrQ = 'narr'
    axes: IterAxesQ = None
    
    def __new__(cls, arr, axes: IterAxesQ = None):
        # Input array is an already formed ndarray instance
        # We first cast to be our class type
        obj = np.asarray(arr).view(cls)
        # add the new attribute to the created instance
        if isinstance(axes, naxes):
            obj.axes = safecopy(axes)
        elif axes is not None:
            axes = naxes(*axes)
        obj.axes = safecopy(axes, cls.baseaxes(obj))
        return obj

    def __getitem__(self, key: AxisKey) -> naxis:   
        return super().__getitem__(key)
        # return np.ndarray.__getitem__(self, key)
                
    def __array_finalize__(self, obj):
        # We do not need to return anything
        if obj is None: return

        axes = getattr(obj, 'axes', self.baseaxes())
        self.axes = axes
        # Ensure correct shape
        if self.ndim != len(self.axes):            
            raise DimsError(self.axes, self.shape)

    def __init__(self, *args, **kwargs):
        # super().__init__()
        self.name = kwargs.get('name', self.name)


                        
    def __array_finalize__(self, obj):     
        # We do not need to return anything
        if obj is None: return

        axes = getattr(obj, 'axes', self.baseaxes())
        self.axes = axes
        # Ensure correct shape
        if self.ndim != len(self.axes):            
            raise DimsError(self.axes, self.shape)
                    
    @overload
    def transpose(self,  axes: Tuple[AxisKey, ...]) -> 'narr': ...
    def transpose(self, *axes: AxisKey):
        # dims = self.axes.copy()
        dims = safecopy(self.axes)
        if len(axes) == 0:
            # If no arguments are provided, reverse the order of axes
            axes = dims.names[::-1]
            
        # If the first argument is a tuple, use it as the axes
        elif len(axes) == 1 and isinstance(axes[0], tuple):
            axes = axes[0]
            
        # Get the order of axes indices  
        # print(dims, axes, self.axes, type(self.axes[0]))
        new_idxs = [self.aidx(ax) for ax in axes]
        # Transpose the NamedAxes
        dims = self.__taxes__(*axes)
        # Transpose the underlying numpy array
        transposed_array = np.transpose(np.asarray(self), axes=new_idxs)
        # Create new narr with the transposed ndarray and the updated NamedAxes
        return self.__class__(transposed_array, dims)
    
    def __repr__(self):
        base = np.ndarray.__repr__(self)
        line_0 = first_line(base)
        spaces = first_digit_loc(line_0)
        spaces = ' ' * (spaces - 1)
        return f'{base}\n{spaces}{self.astrs}'
    
    @property
    def T(self):
        return self.transpose(*self.anames[::-1])
        # return self.transpose(*self.axes[::-1])
