"""
File Purpose: LoiLoader

LOI stands for "Line-Of-Integration", or if you prefer,
    "Lots Of Interpolation", "Let's dO Interpolation", or "Let's do Our Interpolation".

It is useful for loading quantities when interpolating onto a new set of coordinates,
    e.g. a path (or multiple paths) through a box, then maybe integrating along the path.
"""

from .vector_arithmetic import VectorArithmeticLoader
from ...errors import (
    InputError, InputMissingError, InputConflictError,
    FormulaMissingError,
    ComponentValueError,
)
from ...tools import (
    UNSET,
    alias, simple_property, format_docstring,
    xarray_d_grid,
)


class LoiLoader(VectorArithmeticLoader):
    '''Line-Of-Integration (LOI) loader, for interpolating onto a new set of coordinates,
    e.g. a path (or multiple paths) through a box, then maybe integrating along the path.
    '''
    cls_behavior_attrs.register('loi', default='los')
    loi = simple_property('_loi', default='los',
        doc='''Line Of Integration (dim, dict of arrays for interp), or "Lots Of Interpolation", if you prefer.
        Used by self("var_loi") pattern (for the interpolation) and any quantities which integrate along a line.
        ([TODO] move the integrate-along-a-line patterns here, e.g. from mhd_radiative_loader.)

        Can be any of the following:
            None --> no integrate-along-a-line (those methods will crash), and no interpolation.
                self("var_loi") == self("var").
            dict --> no integrate-along-a-line (those methods will crash), but can interpolate onto these arrays.
                self("var_loi") == self("var").interp(**self.loi)
            str --> special option (e.g. 'los' for "use behavior implied by self.los")
                see self.LOI_OPTIONS for all special options.
            2-tuple --> (integration instructions, interpolation instructions)
                integration instructions can be None (as above) or str telling dim to integrate along.
                    Note: if dim ends with '_dim', assumes dim[:-4] tells coord values,
                        e.g. x_dim implies integration along 'x_dim' but dx comes from 'x' coord,
                        which is a useful pattern if 'x' coord depends on multiple dims.
                    Note: if dim starts with '-', assumes loi_cumsum should be evaluated in reverse order.
                        Physically, this can affect things like optical depth.
                interpolation instructions can be None (as above) or dict (as above).

        (To supply kwargs to the interpolator, provide them in the dict.
            E.g.: dict(x=array(...), y=array(...), assume_sorted=True).
            Or, for clarity: dict(coords=(x=array(...), y=array(...)), assume_sorted=True).)''')

    LOI_OPTIONS = {
        'los': '''behavior implied by self.los: no interpolation; integration is along line-of-sight dimension.
            Integration only works if line-of-sight implies a dimension (e.g. 'x')
                (line-of-sight can generally be an array).''',
    }

    loi_explicit = alias('loi',
        doc='''loi as 2-tuple of (integration axis, None or dict of arrays for interpolation).
        Getting loi_explicit will crash if no implied integration axis.
        If integration axis starts with '-', it implies loi_cumsum() in reverse order.
            e.g. '-x' is just like 'x' but any loi cumsums will be in reverse.
            Physically, this can affect quantities like optical depth.''')
    @loi_explicit.getter
    def loi_explicit(self):
        idim = self._loi_integration_dim(signed=True)  # <-- crashes if appropriate
        interp = self._loi_interp_dict()  # <-- crashes if appropriate
        return idim, interp

    def _loi_interp_dict(self):
        '''return dict of arrays to interpolate onto, as implied by self.loi.
        result is None if self.loi implies "no interpolation".
        Crash if invalid self.loi.
        '''
        loi = self.loi
        if loi is None:
            result = None
        elif isinstance(loi, dict):
            result = loi
        elif isinstance(loi, str):
            if loi == 'los':
                result = None
            elif loi in self.LOI_OPTIONS:
                raise NotImplementedError(f'Valid LOI string option {loi!r}; subclass should implement.')
            else:
                raise InputError(f'Invalid LOI string option {loi!r}. Options: {list(self.LOI_OPTIONS.keys())}.')
        elif len(loi) == 2:
            result = loi[1]
        else:
            raise InputError(f'Invalid LOI value {loi!r}. See help(type(self).loi) for valid options.')
        return result

    def _loi_integration_dim(self, *, signed=False):
        '''return dimension along which to integrate, as implied by self.loi.
        Crash with InputMissingError if self.loi implies no integration.

        signed: bool
            whether to start with '-' to imply "cumsum in reverse order" inside loi_cumsum().
            E.g. '-z' --> array.isel(z=slice(None,None,-1)).cumsum('z').isel(z=slice(None,None,-1)
        '''
        loi = self.loi
        neg = False
        if loi is None or isinstance(loi, dict):
            raise InputMissingError(f'self.loi={loi!r} implies no integration; cannot get integration dim.')
        elif isinstance(loi, str):
            if loi == 'los':
                los = self.los
                if not isinstance(los, str):
                    raise InputConflictError(f'self.loi="los" requires string los, but got self.los={self.los!r}')
                if los.startswith('-'):
                    neg = True
                    los = los[1:]
                if los in ('x', 'y', 'z'):
                    result = los
                elif los == 'component':
                    if self.current_n_component() == 1:
                        result = str(self.component)
                    else:
                        errmsg = (f'When LOI option "los" implies integration along self.los == {self.los!r}, '
                                  f'need self.current_n_component()==1, but got self.component={self.component!r}.')
                        raise ComponentValueError(errmsg)
                else:
                    raise NotImplementedError(f'self.loi="los" when self.los={self.los!r}')
            else:
                raise InputError(f'Invalid LOI string option {loi!r}. Options: {list(self.LOI_OPTIONS.keys())}.')
        elif len(loi) == 2:
            if loi[0] is None:
                raise InputMissingError(f'self.loi={loi!r} implies no integration; cannot get integration dim.')
            elif isinstance(loi[0], str):
                if loi[0].startswith('-'):
                    neg = True
                    result = loi[0][1:]
            else:
                raise InputError(f'Invalid self.loi[0]. Expected None or str, got {loi[0]!r}.')
        else:
            raise InputError(f'Invalid LOI value {loi!r}. See help(type(self).loi) for valid options.')
        if signed and neg:
            result = '-' + result
        return result

    def _loi_integration_coord(self):
        '''return coord to use for integration, as implied by self.loi.
        Just like self._loi_integration_dim(), but if result ends with '_dim', remove it.
        E.g. if _loi_integration_dim() == 'grid_dim', this returns 'grid'.
        '''
        idim = self._loi_integration_dim()
        return idim[:-len('_dim')] if idim.endswith('_dim') else idim
    
    @known_pattern(r'(.+)_loi', deps=[0])
    def get_var_loi(self, var, *, _match=None):
        '''var_loi = var interpolated onto coordinates implied by self.loi.
        Internally calls self.loi_interp(val) to do the interpolation (if any);
            can reliably use that method for same functionality as array_loi pattern, anywhere.

        Like all kwarg/attrs, can set loi beforehand, or as kwarg inside the call:
            self.loi = value; self('u_loi')  # option 1
            self('u_loi', loi=value)  # option 2
        '''
        here, = _match.groups()
        arr = self(here)
        return self.loi_interp(arr)

    def loi_interp(self, array):
        '''return array interpolated onto coordinates implied by self.loi.
        See help(type(self).loi) for more details.
        Returns array unchanged if self.loi implies no interpolation.
        '''
        loi_interp_dict = self._loi_interp_dict()
        if loi_interp_dict is None:
            result = array
        else:
            result = array.interp(**loi_interp_dict)
        return result

    @known_pattern(r'loi_cumsum_(.+)', deps=[0])
    def get_loi_cumsum_var(self, var, *, _match=None):
        '''loi_cumsum_var = var interpolated onto coordinates implied by self.loi,
        then cumulatively summed along line-of-integration dimension implied by self.loi.
        Internally calls self.loi_cumsum(val); can reliably use that method instead if desired.
        '''
        here, = _match.groups()
        arr = self(here)
        return self.loi_cumsum(arr, do_interp=True)

    cls_behavior_attrs.register('loi_cumsum_sort', default=True)
    loi_cumsum_sort = simple_property('_loi_cumsum_sort', default=True,
        doc='''Tells whether to sort along the integration dimension before doing cumsum in self.loi_cumsum().
        (If sorting changes order, will not restore original order after cumsum operation,
            so operations relying on xarray coord matching would still work,
            but underlying data order will be different between result and original array.)
        True or 'ascending' --> sort in ascending order.
        'descending' --> sort in descending order.
        False --> do not sort.
        This check occurs before any reversing implied by self.loi.
        Example: self.loi_explicit[0] == '-z', loi_cumsum_sort=True
            --> array.sortby('z').isel(z=slice(None,None,-1)).cumsum('z').isel(z=slice(None,None,-1))''')

    @format_docstring(_sort_doc=loi_cumsum_sort.__doc__, sub_ntab=1)
    def loi_cumsum(self, array, *, do_interp=True, sort=UNSET):
        '''return array interpolated onto coordinates implied by self.loi,
        then cumulatively summed along line-of-integration dimension implied by self.loi.
        Cumulative sum is in forward order by default, but if loi integration dim starts with '-',
            instead reverse order, then cumsum, then reverse again to restore original order.

        do_interp: bool
            whether to do the interpolation step (if any) implied by self.loi.
            (disable if providing array which is already the result of loi_interp().)
        sort: UNSET, bool, 'ascending', or 'descending'
            if UNSET, use sort = self.loi_cumsum_sort
            {_sort_doc}

        See help(type(self).loi) for more details.
        '''
        if do_interp:
            array = self.loi_interp(array)
        idim = self._loi_integration_dim(signed=True)
        neg = idim.startswith('-')
        if neg:
            idim = idim[1:]
        if sort is UNSET:
            sort = self.loi_cumsum_sort
        if sort in (True, 'ascending'):
            array = array.sortby(idim)
        elif sort == 'descending':
            array = array.sortby(idim, ascending=False)
        if neg:
            reverse = {idim: slice(None,None,-1)}
            result = array.isel(reverse).cumsum(idim).isel(reverse)
        else:
            result = array.cumsum(idim)
        return result

    # [TODO] integrate-along-a-line patterns could go here.

    cls_behavior_attrs.register('loi_ustr', default=UNSET)
    loi_ustr = simple_property('_loi_ustr', default=UNSET,
        doc='''units string (e.g., 'length') for loi interped coords,
        in case self.units != self.coords_units, and loi not just along a maindim.

        Probably want to use 'length', but 'length' not assumed by LoiLoader,
            because maindims can technically represent anything, not just physical length.
        Subclass can alter default value used by editing _LOI_USTR_DEFAULT.
            (So, if subclass knows maindims and loi are lengths, it should set _LOI_USTR_DEFAULT='length')

        UNSET --> internally uses self._loi_ustr_default (default in LoiLoader: None. Subclass may override.)
        None --> if any units ambiguity (e.g. when getting dloi), crash.
        str --> use this string to get unit conversion factors as needed.
            e.g., umul = self.u(self.loi_ustr, units=self.units, convert_from=self.coords_units_explicit)
                dloi = xarray_d_grid(interped, xc) * umul   # (when loi not in maindims)''')

    loi_ustr_explicit = alias('loi_ustr',
        doc='''loi_ustr, but returns self._loi_ustr_default if loi_ustr is UNSET''')
    @loi_ustr_explicit.getter
    def loi_ustr_explicit(self):
        return self._loi_ustr_default if self.loi_ustr is UNSET else self.loi_ustr

    _loi_ustr_default = None  # subclass can override if it knows what the default should be, e.g. 'length'.

    @known_var
    def get_dloi(self):
        '''dloi = coord spacing to use for integration as implied by self.loi.
        E.g. if self.loi implies integration along 'x' then this is basically dx
            (accounting for self.slices as needed; see also: self.get_maindims_dim_didx)

        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.)

        See also: xarray_d_grid
        '''
        x = self._loi_integration_dim()  # e.g. 'x' if 'component' mode and self.component=='x'.
        # (also, e.g., 'ray_s_dim' or 'ray_s' if in 'ray_draping' mode.)
        xc = self._loi_integration_coord()  # e.g. 'x' if 'x' above
        # (also, e.g., 'ray_s' if in 'ray_draping' mode.)

        if not hasattr(self, 'maindims'):
            raise NotImplementedError('[TODO] get_dloi() for self without maindims')
            # ^ theoretically should be allowed, as long as d{x} is defined,
            #  for the appropriate x, for the relevant subclass...
            #  and, e.g., might be possible if refactoring below,
            #    to avoid relying on zeros_like_maindims().
            #  This might become necessary e.g. for FromDatasetCalculator.
        # else, self.maindims exists
        if x in self.maindims:
            result = self(f'maindims_{x}_dslice')
        else:
            if self.units != self.coords_units_explicit and (self.loi_ustr is None):
                errmsg = (f'self.loi implies integration along {x} coord,\n'
                          f'but self.coords_units ({self.coords_units!r}) != self.units ({self.units!r}),\n'
                          'and self.loi_ustr is None, so cannot reliably get dloi with proper units.\n'
                          "Consider providing self.loi_ustr, e.g. self.loi_ustr='length'.")
                raise FormulaMissingError(errmsg)
            zeros = self.zeros_like_maindims() 
            interped = self.loi_interp(zeros)
            # [TODO][EFF] maybe more efficient to not fully interpolate the array;
            #  might rewrite if this part is a bottleneck. But, in the meantime,
            #  it should give correct answer, so it is probably good enough for now.
            result = xarray_d_grid(interped, xc)
            if self.units != self.coords_units_explicit:
                umul = self.u(self.loi_ustr_explicit, units=self.units, convert_from=self.coords_units_explicit)
                # e.g. umul = self.u('length', 'cgs', 'si') == 100, when user wants cgs values but si coord units.
                result = result * umul
        return result
