"""
File Purpose: RayDrapingLoader

This provides "field-of-view" grid of rays through a box draped across a sphere.
Intended for use with calculators having maindims = ('x', 'y', 'z').

For more details see tools.sci_tools.ray_draping.
"""

import numpy as np

from .patterns import VectorArithmeticLoader, LoiLoader
from .quantity_loader import QuantityLoader

from ..errors import (
    InputMissingError, InputError,
    DimensionalityError, DimensionValueError,
)
from ..tools import (
    alias, simple_property, format_docstring,
    RayDraping,
    xarray_take_xyz,
)

class RayDrapingLoader(QuantityLoader):
    '''ray draping loader, enabling ability to interpolate to ray paths through box,
    defined by field-of-view grid of rays intersecting box region draped across a sphere.

    For more details see help(self.ray_draping_cls)
    '''

    # # # RAY_DRAPER # # #

    ray_draping_cls = RayDraping

    cls_behavior_attrs.register('ray_draping', default=None)
    ray_draping = simple_property('_ray_draping', default=None,
        doc='''None, dict, or RayDraping instance defining RayDraping for ray draping problems.
        Set kwargs here, then get the RayDraping instance via self.ray_draper. E.g.:
            self.ray_draping=dict(...)  # <-- set RayDraping kwargs as desired
            self.ray_draper.rays(system='box')   # <-- ray paths in "box" system

        None --> quantities needing ray_draping will crash.
        dict --> will use RayDraping.from_dict({**self.ray_draping_defaults(), **ray_draping})
                (e.g., XEXT, YEXT, and ZEXT will come from self.maindims_coords() by default.)
        RayDraping --> will use that RayDraping instance directly.''')

    ray_draper = alias('ray_draping_explicit',
        doc='''The RayDraping instance to use. alias to self.ray_draping_explicit.
        Use self.ray_draper to access the RayDraping instance implied by self.ray_draping,
            regardless of whether it is None, dict, or already a RayDraping instance.
        See help(type(self).ray_draping) for more details.''')

    ray_draping_explicit = alias('ray_draping',
        doc='''self.ray_draping, but convert to RayDraping instance if needed.
        (crashes with InputMissingError if self.ray_draping is None.)''')
    @ray_draping_explicit.getter
    def ray_draping_explicit(self):
        rd = self.ray_draping
        if rd is None:
            errmsg = ('RayDrapingLoader cannot determine RayDraping configuration when ray_draping is None.\n'
                      'Consider setting self.ray_draping = dict(...) or RayDraping(...).\n'
                      '(See help(type(self).ray_draping) for details.)')
            raise InputMissingError(errmsg)
        elif isinstance(rd, RayDraping):
            return rd
        elif isinstance(rd, dict):
            kws = {**self.ray_draping_defaults(), **rd}
            return self.ray_draping_cls.from_dict(kws)
        else:
            errmsg = (f"{type(self).__name__}.ray_draping_explicit expects ray_draping to be"
                      f"None, dict, or RayDraping instance; got object with type={type(rd).__name__}.")
            raise InputError(errmsg)

    def ray_draping_defaults(self):
        '''return dict of defaults to use for ray_draping.
        Asserts that set(self.maindims) == {'x', 'y', 'z'},
            and that self.maindims_coords()['z'] >= 0 everywhere.
            (The latter assumption may be lifted eventually,
                but present RayDraping implementation assumes z=0 at bottom.)

        XEXT, YEXT, ZEXT, X0, Y0 will be taken from self.maindims_coords()
        No other keys included here.
        '''
        assert hasattr(self, 'maindims'), "RayDrapingLoader expects to be inherited by maindims haver."
        # ^^ slightly hacky, but including MainDimensionsHaver in RayDrapingLoader bases
        #       explicitly instead led to hard-to-debug inheritance mro() issue with PlasmaCalculator.
        if set(self.maindims) != {'x', 'y', 'z'}:
            errmsg = (f"{type(self).__name__}.ray_draping_defaults() expects maindims=={{'x','y','z'}} "
                      f"but got maindims={self.maindims}.\nRay draping on non-xyz dims is not yet implemented.\n"
                      "Workaround option: explicitly set ray_draping = RayDraping(...)")
            raise DimensionalityError(errmsg)
        coords = self.get_maindims_coords()
        xmin, xmax = coords['x'].min(), coords['x'].max()
        ymin, ymax = coords['y'].min(), coords['y'].max()
        zmin, zmax = coords['z'].min(), coords['z'].max()
        xext = xmax - xmin
        yext = ymax - ymin
        zext = zmax - zmin
        if coords['z'].min() != 0:
            errmsg = (f"{type(self).__name__}.ray_draping_defaults() expects min(z)==0,\n"
                      f"but got z with min(z)={coords['z'].min()}.\n"
                      "Ray draping with min(z) != 0 is not yet implemented. "
                      "Workaround option: explicitly set ray_draping = RayDraping(...)")
            raise DimensionValueError(errmsg)
        return {'XEXT': xext, 'YEXT': yext, 'ZEXT': zext, 'X0': xmin, 'Y0': ymin}


    # # # LOS & LOI OPTIONS # # #

    LOS_OPTIONS = {**VectorArithmeticLoader.LOS_OPTIONS,
        'ray_draping': '''var_los == self("var_los", los=self.ray_draper.ray_tangents()).
            i.e., use line-of-sight along rays implied by self.ray_draping.''',
        '-ray_draping': '''var_los == self("var_los", los=-self.ray_draper.ray_tangents()).
            i.e., antiparallel to line-of-sight along rays implied by self.ray_draping''',
    }

    @known_var
    @format_docstring(_doc_=VectorArithmeticLoader.get_los.__doc__)
    def get_los(self):
        '''{_doc_}'''
        los = self.los
        if los in ('ray_draping', '-ray_draping'):
            draper = self.ray_draper
            ray_tangents = draper.ray_tangents()
            if los == 'ray_draping':
                result = ray_tangents
            else:
                result = -ray_tangents
        else:
            result = super().get_los()
        return result

    LOI_OPTIONS = {**LoiLoader.LOI_OPTIONS,
        'los': '''behavior implied by self.los. Probably no interpolation, integration along los dimension.
            However, one special case is if los=='ray_draping' or '-ray_draping'.
                In that case, this is equivalent to using self.loi = 'ray_draping'.
            Otherwise, no interpolation, and integration only works if los implies a dim (e.g. 'x')
                (line-of-sight can generally be an array).''',
        'ray_draping': '''var_loi = self("var_loi", loi=(None, self.ray_draper.rays().pc.take_xyz(as_dict=True))).
            i.e., interpolate to loi ray paths implied by self.ray_draping.
            Meanwhile, integration is along the ray dimension ('ray_s' or 'ray_s_dim').''',
        # [TODO] does loi=='ray_draping' actually work? Or does it cause crash when self.los isn't also 'ray_draping'?
    }

    def _loi_is_ray_draping(self):
        '''returns whether self.loi implies to do ray draping.
        True if self.loi == 'ray_draping', or if self.loi == 'los' and los == 'ray_draping' or '-ray_draping'.
        '''
        if self.loi == 'los':
            return isinstance(self.los, str) and self.los in ('ray_draping', '-ray_draping')
        else:
            return self.loi == 'ray_draping'

    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.
        '''
        if self._loi_is_ray_draping():
            result = xarray_take_xyz(self.ray_draper.rays('box'), as_dict=True)
        else:
            result = super()._loi_interp_dict()
        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') instead of array.cumsum('z').
        '''
        if self._loi_is_ray_draping():
            result = self.ray_draper.ray_s_dim
            neg = self.los.startswith('-')
            if signed and neg:
                result = '-' + result
        else:
            result = super()._loi_integration_dim(signed=signed)
        return result

    @property
    def _loi_ustr_default(self):
        '''default unit string to use for loi-related units, e.g. when calculating dloi.
        here, returns 'length' if self._loi_is_ray_draping(), else result from super().
        '''
        return 'length' if self._loi_is_ray_draping() else super()._loi_ustr_default
