"""
File Purpose MaindimsQuantitiesLoader
loading quantities related to the main dimensions
"""

from .quantity_loader import QuantityLoader
from ..dimensions import MainDimensionsHaver
from ..tools import xarray_np_gradient

class MaindimsQuantitiesLoader(MainDimensionsHaver, QuantityLoader):
    '''loader for quantities related to the main dimensions.'''

    @known_var
    def get_maindims_0(self):
        '''return array of 0s, with dimensions and coords implied by self.maindims.
        To save space, the result's dtype is int, by default.
        Equivalent: self.zeros_like_maindims()
        '''
        return self.zeros_like_maindims()

    @known_pattern(r'maindims_(.+)_coord')
    def get_maindims_dim_coord(self, var, *, _match=None):
        '''maindims_x_coord = array of x values and x coordinates.
        Equivalent: self.get_maindims_coord_xarrays()[x], self('maindims_0')[x]
        x must be one of self.maindims.

        Result's values use self.units; coords use self.coords_units (default==self.units).
            (e.g., units='cgs', coord_units='si' --> values are in cgs, but coords are in SI.)
        '''
        x, = _match.groups()
        self._assert_is_maindim(x)
        return self.get_maindims_coord_xarrays()[x]

    @known_pattern(r'maindims_(.+)_idx')
    def get_maindims_dim_idx(self, var, *, _match=None):
        '''maindims_x_idx == array of index associated with x values,
        based on self.slices. Only works if x in self.maindims.
        E.g. if slices[x]=slice(10, None, 4), then result values = [10, 14, 18, ...].
        Equivalent: self.get_maindims_coords_idx['x']
        '''
        x, = _match.groups()
        self._assert_is_maindim(x)
        return self.get_maindims_coords_idx()[x]

    @known_pattern(r'maindims_(.+)_didx', deps=[{0: 'maindims_{group0}_idx'}])
    def get_maindims_dim_didx(self, var, *, _match=None):
        '''maindims_x_didx == array of index spacing associated with x values,
        based on self.slices. Only works if x in self.maindims.
        E.g. if slices[x]=slice(10, None, 4), then result values = [4, 4, 4, ...].
        Equivalent: xarray_np_gradient(self('maindims_{x}_idx'), x)
        '''
        x, = _match.groups()
        idx = self(f'maindims_{x}_idx')  # asserts x is maindim
        return xarray_np_gradient(idx, x)

    @known_pattern(r'maindims_(.+)_dgrid')
    def get_maindims_dim_dgrid(self, var, *, _match=None):
        '''maindims_x_dgrid == array of underlying grid spacing associated with x values.
        Result's values use self.units; coords use self.coords_units (default==self.units).
            (e.g., units='cgs', coord_units='si' --> values are in cgs, but coords are in SI.)

        Caution: this is the underlying grid cell size, regardless of self.slices;
            Thus, this value is suitable for asking questions about the underlying grid,
                e.g. "what is the smallest wavelength the grid can support?"
            But it is NOT suitable for doing an integral across this dim, because, e.g.,
                setting self.slices[x]=slice(None,None,4) does NOT lead to 4 times larger result.
            For doing an integral, consider instead: self('maindims_{x}_dslice')

        Equivalent: self.get_maindims_dgrid(x)
        '''
        x, = _match.groups()
        self._assert_is_maindim(x)
        return self.get_maindims_dgrid(x)

    @known_pattern(r'maindims_(.+)_dslice',
                   deps=[{0: 'maindims_{group0}_dgrid'}, {0: 'maindims_{group0}_didx'}])
    def get_maindims_dim_dslice(self, var, *, _match=None):
        '''maindims_x_dslice == array of spacing between adjacent x values, accounting for self.slices.
        Result's values use self.units; coords use self.coords_units (default==self.units).
            (e.g., units='cgs', coord_units='si' --> values are in cgs, but coords are in SI.)

        Thus, this value is suitable for doing an integral.
            E.g., self.slices[x]=slice(None,None,4) should make ~4 times larger result than slice(None).
                (only "roughly", because the gradient accounts for adjacent values in the dgrid array.)
            For learning about the underlying grid, consider self('maindims_{x}_dgrid') instead.

        Equivalent: maindims_x_dgrid * maindims_x_didx.
        '''
        x, = _match.groups()
        dgrid = self(f'maindims_{x}_dgrid')
        didx = self(f'maindims_{x}_didx')
        return dgrid * didx
