"""
File Purpose: "field-of-view" grid of rays through a box draped across a sphere.

Useful for making synthetic observations based on a plane-parallel (uncurved) simulation.
The overarching idea is:
    (1) Pick a FOV, like a grid of points in a rectangle on the sky.
    (2) Anchor the simulation box to the surface of a sphere.
    (3) "Drape" the box across the sphere, respecting curvature.
    (4) Trace rays from observer, through FOV, through draped box.

There are two "coordinate" systems (though, chosen coords can vary...) to keep in mind:
    (A) "SPH": describes the situation "as understood by the sphere".
        Rays are all straight lines in this system.
        Best coordinate systems for SPH:
            - spherical coordinates centered at sphere center
            - cartesian coordinates centered anywhere (but probably at sphere center)
    (B) "BOX": describes the situation "as understood by the (uncurved) simulation box".
        Rays are generally curved paths in this system.
        Best coordinate systems for BOX:
            - cartesian coordinates directly from simulation
            - cartesian coordinates centered at anchor point

There are various miscellaneous intracies:
    (1) Choosing FOV is simple enough. Using xarray makes it easy to have a grid of points.
    (2) Choosing anchor is simple enough, just need to pick a point in both coord systems.
        e.g. Abox=(x=X0+XEXT/2, y=Y0+YEXT/2, z=0); Asph=(r=R0, θ=(user choice), φ=(user choice)),
            where R0 = radius of sphere,
            XEXT, YEXT are x, y extents of the box,
            and X0, Y0 are x, y coordinates of lower-left corner of the box, default (0,0).
    (3) Draping the simulation is conceptually challenging but can be figured out,
        and the resulting formulas "aren't too bad". Here's what we want:
        (a) (going up in BOX) == (going same distance radially outward in SPH).
            --> zbox - Abox_z = rsph - Asph_r
            --> zbox = rsph - R0   (assuming Abox_z=0 and Asph_r=R0 as done in this file.)
        (b) (going horizontally in BOX at z=0) == (going same distance along sphere surface)
        Combining (a) and (b) almost fully specifies the mapping constraints,
        but there is one more choice to make: the direction for horizontal "motion".
        Basically, we need to pick two orthogonal unit vectors in the SPH system,
        in the plane tangent to the sphere at Asph ("tangent" ensures no changes in rsph).
        There are many ways to do this, but for this module let's choose:
        (c) (θhat_Asph <--> box x motion); (φhat_Asph <--> box y motion)
            Where unit vectors are defined in the usual way:
            rhat_Asph = sin(Asph_θ) cos(Asph_φ) xhat_sph
                       + sin(Asph_θ) sin(Asph_φ) yhat_sph
                       + cos(Asph_θ) zhat_sph
            θhat_Asph = cos(Asph_θ) cos(Asph_φ) xhat_sph
                       + cos(Asph_θ) sin(Asph_φ) yhat_sph
                       - sin(Asph_θ) zhat_sph
            φhat_Asph = - sin(Asph_φ) xhat_sph
                        + cos(Asph_φ) yhat_sph
        (TODO: allow any e1hat and e2hat for x and y, mixing θhat_Asph and φhat_Asph.)
        Now we have fully specified the mapping constraints.
    (4) For tracing rays, we need to know when the ray overlaps with the "draped box" region.
        Assuming we are allowed to periodically tile the box in x and y (and, around the sphere),
        this is easy enough. For a ray expressed parametrically as w(s) (vector),
        the ray overlaps the draped box region for all s such that:
            R0 <= |w(s)| <= R0 + ZEXT,
        where ZEXT is the extent of the simulation box above z=0.
        (If desired, the lower bound can be decreased to allow boxes which extend "below the surface".)
        This expression might be solved analytically to choose s,
        or it might just be imposed as a condition (e.g., mask values where it is False).

A bit more mapping logic:
    For clarity, let's first define Vsph:
            Vsph(xbox,ybox) = (xbox - Abox_x) θhat_Asph + (ybox - Abox_y) φhat_Asph
        Consider the tangent plane at Asph, before "draping" it onto the surface;
        the point in that plane corresponding to (xbox,ybox) in the simulation is:
            vsph(xbox,ybox,zbox=0) = Asph + Vsph(xbox,ybox)
        This is a conceptually useful definition which will reappear in the formulas.
    The mapping itself can now be determined algebraically by applying Rodrigues rotation formula,
        with angle:
            α = |Vsph(xbox,ybox)|/R0
        around axis of rotation:
            ghat_sph = rhat_Asph x Vsph_hat
            (where Vsph_hat = Vsph/|Vsph|)
        This determines the mapping because of (3b) above; "going a distance along the surface"
        corresponds to "rotating around by the correct amount, around the axis which is perpendicular
        to the plane containing the path as well as the sphere center." That axis is determined by
        rhat_Asph x Vhat_sph because those are two orthogonal vectors which lie in that plane.
    Finding the inverse mapping just involves algebra, starting from the mapping itself.
    
The mapping itself:
    ANCHOR POINT UNIT VECTORS (these are constants):
        rhat_Asph = sin(Asph_θ) cos(Asph_φ) xhat_sph
                   + sin(Asph_θ) sin(Asph_φ) yhat_sph
                   - cos(Asph_θ) zhat_sph
        θhat_Asph = cos(Asph_θ) cos(Asph_φ) xhat_sph
                   + cos(Asph_θ) sin(Asph_φ) yhat_sph
                   - sin(Asph_θ) zhat_sph
        φhat_Asph = - sin(Asph_φ) xhat_sph
                    + cos(Asph_φ) yhat_sph
    BOX TO SPH: (xbox,ybox,zbox) --> wsph (vector):
        rsph = R0 + zbox
        Vsph = (xbox - Abox_x) θhat_Asph + (ybox - Abox_y) φhat_Asph
        Vsph_hat = Vsph/|Vsph|
        α = |Vsph|/R0
        wsph = (cos(α) rhat_Asph + sin(α) Vsph_hat) rsph
    SPH TO BOX: wsph (vector) --> (xbox,ybox,zbox):
        rsph = |wsph|
        |Vsph| = R0 cos^-1((wsph . rhat_Asph)/rsph)
        wsph_perp = wsph - (wsph . rhat_Asph) rhat_Asph == rsph sin(α) Vsph_hat
        Vsph_hat = wsph_perp/|wsph_perp|
        Vsph = |Vsph| Vsph_hat
        xbox = Abox_x + (Vsph . θhat_Asph)
        ybox = Abox_y + (Vsph . φhat_Asph)
        zbox = rsph - R0
    TRICKY ANGLES:
        When wsph_perp --> 0, it is impossible to get Vsph_hat;
            this corresponds to the point xbox = Abox_x, ybox = Abox_y.
            --> we need a test like "if |wsph_perp| < ε, impose Vsph = 0.
                (maybe ε should be some small fraction of simulation grid cell size.)
        If (vsph . rhat_Asph)/rsph > 1, the inverse cosine will fail.
            this shouldn't occur anywhere though, |wsph| = |rsph| by definition,
            and rhat_Asph is a unit vector, so the whole thing should be 1 or less.
            Numerically might be an issue? Can double-check, but low priority.
        No other tricky angles, the rest should be fine!
    RAY FORMULA:
        Rays in sph system are straight lines:
            w(s) = F + rayhat * s
        where:
            rayhat = (F - T)/|F - T|,
            F is the point in FOV corresponding to this ray (there is one ray per FOV point),
            T is the telescope/observer position. (F and T both expressed in sph system.)
        The corresponding paths in box system are curved by the mapping formula above.
        Assuming periodic box tiling is allowed, the rays overlap the draped box region wherever:
            R0 <= |w(s)| <= R0 + ZEXT
        and the periodic box tiling can be accomplished by taking the modulus:
            xbox = xbox % XEXT
            ybox = ybox % YEXT
            zbox = zbox % ZEXT
    RAY BOUNDARIES:
        The values of s where rays enter and exit the draped box region can be computed from:
            w(s) = F + rayhat * s
            rayhat = (F - T)/|F - T|
            R0 <= |w(s)| <= R0 + ZEXT
        For boundary R = R0 or R0 + ZEXT:
            R^2 = F^2 + 2 (F . rayhat) s + s^2
            --> s^2 + 2 (F . rayhat) s + (F^2 - R^2) = 0
            --> s = - (F . rayhat) ± sqrt((F . rayhat)^2 - F^2 + R^2)
        The full ray intersects the boundary whenever the discriminant is non-negative.
        In that case, is important to know which point the ray passes through first,
            i.e. which is closest to T. That point is the "entry".
            If there are two intersections of a boundary, the other is the "exit".
            Thankfully, by definition, the smaller s value is always the entry point.
            (Not necessarily smaller in magnitude; s might be negative. E.g. -7 is smaller than 3.)
        The full decision tree looks like:
            - find lower and upper boundary entry and exit points. (Up to 4 distinct points).
            - if lower entry == lower exit, ray just grazes lower boundary; ignore lower boundary.
            - if upper entry == upper exit, ray just grazes upper boundary; replace ray with NaNs.
            - if no upper entry/exit, ray doesn't overlap sphere; replace ray with NaNs.
            - if no lower entry/exit, ray goes from upper entry to upper exit. (E.g., near disk edge.)
            - else, upper entry & exit, and lower entry & exit are all distinct,
                so ray goes from upper entry to lower entry (e.g., anywhere other than disk edge),
                and the exits both correspond to the draped region on the other side of the sphere.
        NOTE: when actually getting the standard ray_s values, the default is inclusive=(False, False),
            which avoids points "on the edges". This way all ray_s correspond to points INSIDE the box.
            Without doing this, some interp results may give NaNs due to points on edges.
    FOV POSITION AND SIZE:
        The FOV can be any size and be placed anywhere,
            as long as the sphere does not block it from T (not checked directly in this module...),
            and as long as it is a plane (not actually strictly necessary but probably a good idea).
        In practice, it is nice to choose FOV location based on box.
        How big should FOV be? Ideally, big enough to cover box.
        This defines "standard" FOV (assuming T along zsph axis):
            get points on box surfaces in sph system, project onto zsph=constant plane,
            then choose FOV x and y limits as min and max x and y across all projected points.
            Projecting all 8 corners is a good start but insufficient;
                also need to check top surface which can bulge out beyond the corners.
            Internal algorithm checks all 8 corners and a low-res grid (21x21 points) for top surface.
            (Currently, only implemented if T along zsph axis. Eventually, might handle arbitrary T.)
"""
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr

from ..docs_tools import format_docstring
from ..math import is_integer
from ..properties import alias_child, simple_property
from ..xarray_tools import (
    xarray_sel,
    xr1d, xarray_grid,
    _get_basis, xarray_dot_product,
    xarray_magnitude, xarray_unit_vector,
    xarray_take_xyz, xarray_join_xyz,
    xarray_differentiate,
)
from ...errors import (
    InputError, InputMissingError, InputConflictError,
    TooManyPlottablesError,
    DimensionalityError,
)


### --------------------- Helper Classes (part of setup) --------------------- ###

class RayDrapeScene():
    '''representation of some "background" conditions for the ray draping problem.
    (High-level interface will allow user to change any of these,
        but for internal mental model, let's think of them as "probably constants".)
    Safest to convert to meters throughout, but it is probably technically units agnostic.

    R0: scalar
        radius of sphere (onto which z=0 of box will be draped)
    XEXT, YEXT, ZEXT: scalars
        extents of box to be draped
    T: scalar or xr.DataArray with 'component' dimension
        telescope/observer position (in sph system, cartesian coords centered on sphere)
        if scalar, interpret as distance along sph z axis, equivalent to T * pc.ZHAT.
        E.g., T = 1.5e11 * pc.ZHAT   # 1 AU from sphere center along sph z axis.
        (decent choice because it aligns zbox with zsph when viewing face-on / centered.)
    X0, Y0: scalar, default 0
        x, y coordinates of lower-left corner of box in box system
    '''
    _init_kwargs = {'R0', 'XEXT', 'YEXT', 'ZEXT', 'T', 'X0', 'Y0'}

    def __init__(self, *, R0, XEXT, YEXT, ZEXT, T, X0=0, Y0=0):
        if np.size(T) == 1:
            BASIS = _get_basis()
            ZHAT = BASIS.ZHAT
            T = T * ZHAT
        self.R0 = R0
        self.XEXT = XEXT
        self.YEXT = YEXT
        self.ZEXT = ZEXT
        self.T = T
        self.X0 = X0
        self.Y0 = Y0

    # # # DISPLAY # # #
    def _T_display(self):
        '''T in brief human-readable string if easy to do that, else just something reasonable.
        Extra pretty if 0.1 to 1000 AU (assuming SI units).
            (if non-SI units, might show AU accidentally, but this is just display,
            so it won't affect actual computations of results --> it's worth it.)
        '''
        T = self.T
        Tdisp = T.where(lambda arr: arr != 0, drop=True)
        if Tdisp.size == 1:
            comp = Tdisp['component'].values.item()
            zhat = f'{str(comp).upper()}HAT'
            return f"{Tdisp.values.item():.1e} {zhat}"
        elif T.size == 3:
            mag = np.sqrt(np.sum(T.values**2))
            return f"DataArray(vector with mag={mag:.1e})"
        else:
            return f"DataArray(sizes={dict(T.sizes)})"

    def _X0Y0_display(self):
        '''brief human-readable strings for X0 and Y0 excluding defaults (0).'''
        result = []
        for x in ['X0', 'Y0']:
            val = getattr(self, x)
            if val == 0:
                continue
            else:
                if np.size(val) > 1:
                    x_str = f"{type(val).__name__}(sizes={dict(val.sizes)})"
                else:
                    x_str = f"{val:.1e}"
                result.append(f"{x}={x_str}")
        return result

    def _repr_contents(self):
        '''list of contents for self.__repr__'''
        result = [
            f"R0={self.R0:.1e}",
            f"XEXT={self.XEXT:.1e}",
            f"YEXT={self.YEXT:.1e}",
            f"ZEXT={self.ZEXT:.1e}",
            f"T={self._T_display()}",
        ]
        result.extend(self._X0Y0_display())
        return result

    def __repr__(self):
        contents_str = ', '.join(self._repr_contents())
        return f'{type(self).__name__}({contents_str})'


class RayDrapeAnchor():
    '''representation of Anchor point for ray draping problem.
    For many helpful quantities, consider using self.Asph.
    (see help(type(self).Asph) for details.)

    Currently assumes Asph_r=R0, Abox_z=0.

    Asph_theta: scalar, default 0
        Asph_θ, i.e. anchor polar angle θ in sph system, in radians.
        This is the angle from +z to Asph, in the sph system.
        (for T along sph z axis, θ=0 corresponds to center of disk)
    Asph_phi: scalar, default 0
        Asph_φ, i.e. anchor azimuthal angle φ in sph system, in radians.
        This is the angle from +x to [Asph projected into x-y plane], in the sph system.
        (for T along sph z axis and θ=0, φ=0 provides intuitive x-y sph-box alignment,
            while φ=π/2 rotates box 90° clockwise as viewed from +z, thus "swapping" x and y)
    Abox_x: None or scalar
        Abox_x, i.e. anchor x position in box system
        None --> mapping will assume Abox_x = X0 + XEXT/2
    Abox_y: None or scalar
        Abox_y, i.e. anchor y position in box system
        None --> mapping will assume Abox_y = Y0 + YEXT/2
    '''
    # [TODO]: use UNSET for inputs with defaults; internally use class defaults for UNSET inputs.

    _init_kwargs = {'Asph_theta', 'Asph_phi', 'Abox_x', 'Abox_y'}

    def __init__(self, *, Asph_theta=0, Asph_phi=0, Abox_x=None, Abox_y=None):
        self.Asph_theta = Asph_theta
        self.Asph_phi = Asph_phi
        self.Abox_x = Abox_x
        self.Abox_y = Abox_y

    Asph_theta = simple_property('_Asph_theta',
        doc='''Asph_θ, i.e. anchor polar angle θ in sph system, in radians.
        See help(type(self)) for details.''')
    @Asph_theta.setter
    def Asph_theta(self, value):
        self._Asph_theta = value
        self._Astate = 1 + getattr(self, '_Astate', 0)
    
    Asph_phi = simple_property('_Asph_phi',
        doc='''Asph_φ, i.e. anchor azimuthal angle φ in sph system, in radians.
        See help(type(self)) for details.''')
    @Asph_phi.setter
    def Asph_phi(self, value):
        self._Asph_phi = value
        self._Astate = 1 + getattr(self, '_Astate', 0)

    @property
    def Asph(self):
        '''various details related to Asph_theta and Asph_phi.

        self.Asph is a dict containing:
            'θ': self.Asph_theta
            'cosθ': cos(Asph_theta)
            'sinθ': sin(Asph_theta)
            'φ': self.Asph_phi
            'cosφ': cos(Asph_phi)
            'sinφ': sin(Asph_phi)
            'rhat': rhat_Asph as xr.DataArray vector (along 'component' dim),
                    in sph system cartesian coords centered at r=0
            'θhat': θhat_Asph as xr.DataArray vector (along 'component' dim),
                    in sph system cartesian coords centered at r=0
            'φhat': φhat_Asph as xr.DataArray vector (along 'component' dim),
                    in sph system cartesian coords centered at r=0
        as well as alternate spellings of keys, with same values:
            'theta': alias to 'θ'
            'cos_theta': alias to 'cosθ'
            'sin_theta': alias to 'sinθ'
            'phi': alias to 'φ'
            'cos_phi': alias to 'cosφ'
            'sin_phi': alias to 'sinφ'
            'thetahat': alias to 'θhat'
            'phihat': alias to 'φhat'

        result is cached; only recomputed when Asph_theta or Asph_phi change.
        '''
        Astate, Asph = getattr(self, '_Asph_cached', (None, None))
        if Astate != self._Astate:
            Asph = self._Asph_calcs()
            self._Asph_cached = (self._Astate, Asph)
        return Asph

    def _Asph_calcs(self):
        '''compute various details related to Asph_theta and Asph_phi.
        See help(type(self).Asph) for details
        '''
        θ = self.Asph_theta
        φ = self.Asph_phi
        cosθ = np.cos(θ)
        sinθ = np.sin(θ)
        cosφ = np.cos(φ)
        sinφ = np.sin(φ)
        BASIS = _get_basis()
        XHAT = BASIS.XHAT
        YHAT = BASIS.YHAT
        ZHAT = BASIS.ZHAT
        rhat = (sinθ * cosφ) * XHAT + (sinθ * sinφ) * YHAT + (cosθ) * ZHAT
        θhat = (cosθ * cosφ) * XHAT + (cosθ * sinφ) * YHAT - (sinθ) * ZHAT
        φhat = (-sinφ) * XHAT + (cosφ) * YHAT
        result = {  # put all "nicer spellings" first, then all plain spellings.
            'θ': θ, 'cosθ': cosθ, 'sinθ': sinθ,
            'φ': φ, 'cosφ': cosφ, 'sinφ': sinφ,
            'rhat': rhat,
            'θhat': θhat, 'φhat': φhat,
            'theta': θ, 'cos_theta': cosθ, 'sin_theta': sinθ,
            'phi': φ, 'cos_phi': cosφ, 'sin_phi': sinφ,
            'thetahat': θhat, 'phihat': φhat,
        }
        return result

    def _get_basis(self):
        '''gets BasisXYZ object containing x,y,z basis info.
        Here, just returns _get_basis() which provides _BASIS from tools.xarray_tools.xarray_vectors
            (Component objects if imported the rest of PlasmaCalcs, strings if copy-pasted tools elsewhere.)
        '''
        return _get_basis()

    # # # DISPLAY # # #
    def _Asph_display(self):
        '''brief human-readable strings for Asph_θ and Asph_φ'''
        val = self.Asph_theta
        if np.size(val) > 1:
            θ_str = f"{type(val).__name__}(sizes={dict(val.sizes)})"
        else:
            θ_str = f"{np.degrees(val):.3g}°"
        val = self.Asph_phi
        if np.size(val) > 1:
            φ_str = f"{type(val).__name__}(sizes={dict(val.sizes)})"
        else:
            φ_str = f"{np.degrees(val):.3g}°"
        return [f"Asph_theta={θ_str}", f"Asph_phi={φ_str}"]

    def _Abox_display(self):
        '''brief human-readable strings for Abox_x and Abox_y, excluding defaults (None).'''
        result = []
        for Abox_x in ['Abox_x', 'Abox_y']:
            val = getattr(self, Abox_x)
            if val is None:
                continue
            else:
                if np.size(val) > 1:
                    x_str = f"{type(val).__name__}(sizes={dict(val.sizes)})"
                else:
                    x_str = f"{val:.1e}"
                result.append(f"{Abox_x}={x_str}")
        return result

    def _repr_contents(self):
        '''list of contents for self.__repr__'''
        result = self._Asph_display()
        result.extend(self._Abox_display())
        return result

    def __repr__(self):
        contents_str = ', '.join(self._repr_contents())
        return f'{type(self).__name__}({contents_str})'


class RayDrapeFOV():
    '''representation of Field-of-View grid for ray draping problem.

    FOV: xr.DataArray with 'component' dimension (and possibly any other dimensions too)
        grid of points representing "field of view" expressed in sph system,
        with 'component' dimension associated with x, y, z coords of each point.
    '''
    _init_kwargs = {'FOV'}

    def __init__(self, FOV):
        if not hasattr(FOV, 'dims') or 'component' not in FOV.dims:
            raise InputError("FOV must be an xarray object (e.g. DataArray) "
                             "with 'component' dimension associated with x, y, z coords of each point.")
        self.FOV = FOV

    # # # DISPLAY # # #
    def _repr_contents(self):
        '''list of contents for self.__repr__'''
        result = [f"FOV=DataArray{dict(self.FOV.sizes)}"]
        return result

    def __repr__(self):
        contents_str = ', '.join(self._repr_contents())
        return f'{type(self).__name__}({contents_str})'


### --------------------- RayDraping class --------------------- ###

_paramdocs_ray_draping = {
    'scene': '''dict or RayDrapeScene
        scene conditions for ray draping problem.
        dict --> convert to RayDrapeScene.
            Should have keys: R0, XEXT, YEXT, ZEXT, T.''',
    'anchor': '''dict or RayDrapeAnchor
        anchor point for ray draping problem.
        dict --> convert to RayDrapeAnchor.
            May have any of these keys: Asph_theta, Asph_phi, Abox_x, Abox_y.
            (Missing keys use RayDrapeAnchor defaults.)''',
    'fov': '''None, xr.DataArray, or RayDrapeFOV
        field-of-view grid for ray draping problem.
        None --> specify it later, by setting self.fov = value.
        xr.DataArray --> convert to RayDrapeFOV.
            Should have 'component' dimension associated with x, y, z coords of each point,
            in sph system cartesian coords centered at r=0.''',
    'wbox': '''xarray object (e.g. DataArray)
        points in box system cartesian coords,
        with 'component' dimension associated with x, y, z coords of each point.
        (Possibly useful: wbox=xr.concat([xbox, ybox, zbox], dim='component'))''',
    'wsph': '''xarray object (e.g. DataArray)
        points in sph system cartesian coords centered at r=0,
        with 'component' dimension associated with x, y, z coords of each point.
        (Possibly useful: wsph=xr.concat([xsph, ysph, zsph], dim='component'))''',
    'in_box': '''bool, 'lower', or 'upper'
        whether to check that points will fall inside the draped box region,
        masking all other points using NaNs.
        'lower' --> only check lower bound: R0 <= |wsph|.
        'upper' --> only check upper bound: |wsph| <= R0 + ZEXT.''',
    'periodic': '''bool
        whether to wrap points periodically into x and y box extents.
        True --> like draping the box periodically around the sphere.
        False --> mainly useful for debugging... are there other use cases?''',
}


@format_docstring(**_paramdocs_ray_draping)
class RayDrapingCore():
    '''core methods for the entire ray draping problem, including mapping calculations.
    This includes the setup (scene, anchor, fov), and mapping methods (for any points).

    Users should prefer RayDraping instead, which additionally includes
        the actual rays, some plotting methods, and some clever defaults.

    scene: {scene}
    anchor: {anchor}
    fov: {fov}
    '''
    scene_cls = RayDrapeScene
    anchor_cls = RayDrapeAnchor
    fov_cls = RayDrapeFOV

    # treat values as 0 if less than TINYTOL * np.finfo(...).eps,
    #  when doing the "tricky angles" check ("if |wsph_perp| < ε, impose Vsph = 0.")
    _TINYTOL = 1e2

    def __init__(self, *, scene, anchor, fov=None):
        if not isinstance(scene, RayDrapeScene):
            scene = self.scene_cls(**scene)
        if not isinstance(anchor, RayDrapeAnchor):
            anchor = self.anchor_cls(**anchor)
        self.scene = scene
        self.anchor = anchor
        self.fov = fov   # internally class self.fov_cls(fov) if needed; see type(self).fov.


    # # # MISC PROPERTIES # # #

    R0 = alias_child('scene', 'R0')
    XEXT = alias_child('scene', 'XEXT')
    YEXT = alias_child('scene', 'YEXT')
    ZEXT = alias_child('scene', 'ZEXT')
    T = alias_child('scene', 'T')
    X0 = alias_child('scene', 'X0')
    Y0 = alias_child('scene', 'Y0')
    Asph = alias_child('anchor', 'Asph')
    Abox_x = alias_child('anchor', 'Abox_x')
    Abox_y = alias_child('anchor', 'Abox_y')
    FOV = alias_child('fov', 'FOV')

    @Abox_x.getter
    def Abox_x(self):
        v = self.anchor.Abox_x
        return self.X0 + self.XEXT / 2 if v is None else v

    @Abox_y.getter
    def Abox_y(self):
        v = self.anchor.Abox_y
        return self.Y0 + self.YEXT / 2 if v is None else v

    _get_basis = alias_child('anchor', '_get_basis')


    # # # fov management # # #

    fov = simple_property('_fov',
        doc='''field-of-view grid for ray draping problem.
        For efficiency, RayDrapeFOV is created immediately when setting self.fov (unless None).
        (E.g., self.fov=arr; self.fov  # == RayDrapeFOV(arr))
        See help(type(self)) for more details about possible values.''')
    @fov.setter
    def fov(self, value):
        self._fov = self._get_fov_value(value)

    def _get_fov_value(self, value):
        '''convert value to RayDrapeFOV (unless None) appropriate for storage in self.fov.'''
        if value is None:
            return None
        elif not isinstance(value, RayDrapeFOV):
            value = self.fov_cls(value)
        return value

    @FOV.getter
    def FOV(self):
        '''alias to self.fov.FOV, this is the actual xr.DataArray associated with the FOV.
        Will raise InputMissingError if self.fov is None.
        '''
        if self.fov is None:
            raise InputMissingError("cannot get FOV when self.fov is None.")
        return self.fov.FOV


    # # # MAPPING - GENERIC FORMULAS # # #

    @format_docstring(**_paramdocs_ray_draping, sub_ntab=1)
    def box_to_sph(self, wbox):
        '''return wsph corresponding to any wbox point or points.
        i.e., returns w expressed in sph system cartesian coords centered at r=0.

        wbox: {wbox}
        '''
        xbox = xarray_sel(wbox, component='x')
        ybox = xarray_sel(wbox, component='y')
        zbox = xarray_sel(wbox, component='z')
        rsph = self.R0 + zbox
        Vsph = (xbox - self.Abox_x) * self.Asph['θhat'] + (ybox - self.Abox_y) * self.Asph['φhat']
        Vsph_mag = xarray_magnitude(Vsph)
        Vsph_hat = Vsph / Vsph_mag
        # check: "if |wsph_perp| < ε, impose Vsph = 0."
        TINY = np.finfo(Vsph_mag.dtype).eps * self._TINYTOL
        Vsph_hat = Vsph_hat.where(Vsph_mag > TINY, other=0)
        # (resume mapping math stuff)
        α = Vsph_mag / self.R0
        wsph = (np.cos(α) * self.Asph['rhat'] + np.sin(α) * Vsph_hat) * rsph
        return wsph

    @format_docstring(**_paramdocs_ray_draping, sub_ntab=1)
    def sph_to_box(self, wsph, *, in_box=False, periodic=True):
        '''return wbox corresponding to any wsph point or points.
        i.e., returns w expressed in box system cartesian coords.

        wsph: {wsph}
        in_box: {in_box}
            (Use False for efficiency if this was already checked elsewhere!)
        periodic: {periodic}
        '''
        # mapping math
        rsph = xarray_magnitude(wsph)
        wsph_dot_rhat = xarray_dot_product(wsph, self.Asph['rhat'])
        Vsph_mag = self.R0 * np.arccos(wsph_dot_rhat / rsph)
        wsph_perp = wsph - wsph_dot_rhat * self.Asph['rhat']
        wsph_perp_mag = xarray_magnitude(wsph_perp)
        Vsph_hat = wsph_perp / wsph_perp_mag
        Vsph = Vsph_mag * Vsph_hat
        # check: "if |wsph_perp| < ε, impose Vsph = 0."
        TINY = np.finfo(wsph_perp_mag.dtype).eps * self._TINYTOL
        Vsph = Vsph.where(wsph_perp_mag > TINY, other=0)
        # (resume mapping math stuff)
        xbox = self.Abox_x + xarray_dot_product(Vsph, self.Asph['θhat'])
        ybox = self.Abox_y + xarray_dot_product(Vsph, self.Asph['φhat'])
        zbox = rsph - self.R0
        # add component labels and join
        wbox = xarray_join_xyz(xbox, ybox, zbox, _basis=self._get_basis())
        # apply keep_in_box and periodic_wrap if desired
        if in_box:
            wbox = self.keep_in_box(wbox, in_box=in_box)
        if periodic:
            wbox = self.periodic_wrap(wbox)
        return wbox

    @format_docstring(**_paramdocs_ray_draping, sub_ntab=1)
    def keep_in_box(self, wsph, *, in_box=True):
        '''return wsph points masked to only those inside the draped box region.
        i.e., only keep points where R0 <= |wsph| <= R0 + ZEXT.

        wsph: {wsph}
        in_box: {in_box}
        '''
        rsph = xarray_magnitude(wsph)
        if in_box == 'lower':
            keep = (self.R0 <= rsph)
        elif in_box == 'upper':
            keep = (rsph <= self.R0 + self.ZEXT)
        elif in_box:
            keep = (self.R0 <= rsph) & (rsph <= self.R0 + self.ZEXT)
        wsph_kept = wsph.where(keep)
        return wsph_kept

    @format_docstring(**_paramdocs_ray_draping, sub_ntab=1)
    def periodic_wrap(self, wbox):
        '''return wbox points wrapped periodically into x and y box extents.
        The wrapping is done using modulus operation, accounting for X0 and Y0 if nonzero:
            xbox = (xbox - X0) % XEXT + X0
            ybox = (ybox - Y0) % YEXT + Y0

        wbox: {wbox}
        '''
        xbox = xarray_sel(wbox, component='x')
        ybox = xarray_sel(wbox, component='y')
        zbox = xarray_sel(wbox, component='z')
        if self.X0 == 0:  # <-- check not required, but can improve efficiency slightly.
            xbox = xbox % self.XEXT
        else:
            xbox = (xbox - self.X0) % self.XEXT + self.X0
        if self.Y0 == 0:  # <-- check not required, but can improve efficiency slightly.
            ybox = ybox % self.YEXT
        else:
            ybox = (ybox - self.Y0) % self.YEXT + self.Y0
        wbox_wrapped = xr.concat([xbox, ybox, zbox], dim='component')
        return wbox_wrapped


    # # # DISPLAY # # #

    def _repr_contents(self):
        '''list of contents for self.__repr__'''
        result = [
            f"scene={repr(self.scene)}",
            f"anchor={repr(self.anchor)}",
            f"fov={repr(self.fov)}",
        ]
        return result

    def __repr__(self):
        contents_str = ', '.join(self._repr_contents())
        return f'{type(self).__name__}({contents_str})'


_paramdocs_ray_draping.update({
    'ray_s': '''None, int, dict, or xr.DataArray including 'ray_s' or 'ray_s_dim' dimension.
        Distance along each ray, from FOV. Rays are parametrized (in sph system) as:
            w(s) = F + rayhat * s,   where:
            rayhat = (F - T)/|F - T|,
            F is the point in FOV corresponding to this ray (there is one ray per FOV point),
            T is the telescope/observer position.
        Must include the 'ray_s' or 'ray_s_dim' dimension. Can be 1D if all rays have same spacing,
            but might also have same dims as FOV, to specify different spacing for each ray.
        None --> specify later (getting self.ray_s before specifying will just cause crash).
        int --> number of points along each ray, choosing self.standard_ray_s_limits() for min and max s,
            while using inclusive=(False,False) when generating the grid (to avoid including edge points).
            standard_ray_s_limits() picks s limits to precisely cover full paths through the draped box region.
        dict --> kwargs for pc.xarray_grid(). Uses self.standard_ray_s_limits() for min and max, unless provided.
            Also uses name='ray_s' unless provided, and inclusive=(False,False) unless provided.
            [TODO] add "maxstep" option to xarray_grid(), to enable "maxstep = simulation dx".
            Workaround to solve that problem for now: specify ray_s=None when making drape, then later do:
                smin, smax = drape.standard_ray_s_limits()
                ray_s=pc.xarray_grid(smin, smax, N=1+int(np.max((smax-smin)/dx)))
        Examples: 100; dict(N=100); pc.xr1d([100, 200, 300, 400], name='ray_s').
        (Use self.ray_s to see ray_s value, and self.ray_s_array to see the actual ray_s xr.DataArray.)''',
})

@format_docstring(**_paramdocs_ray_draping)
class RayDraping(RayDrapingCore):
    '''representation of the entire ray draping problem, including mapping calculations.
    Consider using self.wbox(), to get "the answer", i.e. ray points mapped into box system.

    scene: {scene}
    anchor: {anchor}
    fov: 3-tuple of ("standard", Nx, Ny), {fov}
        ("standard", Nx, Ny) --> use self.standard_fov(Nx, Ny) to create reasonable FOV,
            chosen such as the rectangle at zsph=R0+ZEXT) which fully covers entire box
            (without any smart rotations yet; edges parallel to x sph and y sph axes).
    ray_s: {ray_s}


    EXAMPLE
    -------
        # Below is a somewhat exaggerated example for visualization purposes;
        #   if the 3D simulation box is really that large compared to the Sun,
        #   then curvature effects might be physically significant within the plasma,
        #   so maybe a plane-parallel simulation isn't the best idea on that scale.
        import PlasmaCalcs as pc
        import numpy as np
        Rsun = 6.957e8  # [m]
        AU = pc.DEFAULTS.PHYSICAL.CONSTANTS_SI['AU']   # 1.495978707e11  [m]
        scene = pc.RayDrapeScene(R0=Rsun,
                                 XEXT=1.0*Rsun,
                                 YEXT=0.7*Rsun,
                                 ZEXT=1.0*Rsun,
                                 T=1*AU)
        anchor = pc.RayDrapeAnchor(Asph_theta=np.pi*3/8,
                                   Asph_phi=np.pi/4)
        drape = pc.RayDraping(scene=scene, anchor=anchor,
                              fov=("standard", 7, 4),
                              ray_s=30)
        drape.plot('both', periodic=False)
    '''

    def __init__(self, *, scene, anchor, fov=None, ray_s=None):
        super().__init__(scene=scene, anchor=anchor, fov=fov)
        self.ray_s = ray_s


    @classmethod
    @format_docstring(scene_kw=RayDrapeScene._init_kwargs,
                      anchor_kw=RayDrapeAnchor._init_kwargs,
                      fov_kw=RayDrapeFOV._init_kwargs)
    def from_dict(cls, d):
        '''create RayDraping instance from dict, with keys corresponding to __init__ kwargs,
        or to kwargs used to initialize scene, anchor, and fov, e.g. can provide "T" here.

        Cannot overlap __init__ and scene/anchor/fov kwargs;
            e.g. cannot provide both "T" and "scene".

        Keys allowed are:
            'scene' or scene kwargs: {scene_kw}
            'anchor' or anchor kwargs: {anchor_kw}
            'fov' or fov kwargs: {fov_kw}
            'ray_s'
        '''
        scene_kw = cls.scene_cls._init_kwargs
        anchor_kw = cls.anchor_cls._init_kwargs
        fov_kw = cls.fov_cls._init_kwargs
        # checks:
        if 'scene' in d and any(k in d for k in scene_kw):
            overlap = tuple(k for k in scene_kw if k in d)
            raise InputConflictError(f"Cannot provide both 'scene' and scene kwargs {overlap}.")
        if 'anchor' in d and any(k in d for k in anchor_kw):
            overlap = tuple(k for k in anchor_kw if k in d)
            raise InputConflictError(f"Cannot provide both 'anchor' and anchor kwargs {overlap}.")
        if 'fov' in d and any(k in d for k in fov_kw):
            overlap = tuple(k for k in fov_kw if k in d)
            raise InputConflictError(f"Cannot provide both 'fov' and fov kwargs {overlap}.")
        # method:
        d = {**d}  # copy to avoid modifying input
        scene = d.pop('scene', None)
        if scene is None:
            scene = {k: d.pop(k) for k in scene_kw if k in d}
            scene = cls.scene_cls(**scene)
        anchor = d.pop('anchor', None)
        if anchor is None:
            anchor = {k: d.pop(k) for k in anchor_kw if k in d}
            anchor = cls.anchor_cls(**anchor)
        fov = d.pop('fov', None)
        if fov is None:
            fov = {k: d.pop(k) for k in fov_kw if k in d}
            if fov:
                fov = cls.fov_cls(**fov)
            else:
                fov = None  # if no fov kwargs provided, use None (i.e. specify later)
        ray_s = d.pop('ray_s', None)
        if len(d) > 0:
            raise InputError(f"Unrecognized keys in input dict: {list(d.keys())}")
        return cls(scene=scene, anchor=anchor, fov=fov, ray_s=ray_s)


    # # # MISC PROPERTIES # # #

    @property
    def n_rays(self):
        '''number of rays, i.e. number of points in self.FOV'''
        sizes = self.FOV.sizes
        n_rays = 1
        for dim in sizes:
            if dim != 'component' and dim != 'ray_s' and dim != 'ray_s_dim':
                n_rays = n_rays * sizes[dim]
        return n_rays


    # # # MAPPING - THESE PARTICULAR RAY_S # # #

    def rayhat(self):
        '''rayhat, i.e. unit vector along each ray in sph system cartesian coords.
        
        rayhat has dimensions of self.FOV,
            as well as 'component' dimension associated with x, y, z coords of each point.

        rayhat = (F - T)/|F - T|,   where:
            F is the point in FOV corresponding to this ray (there is one ray per FOV point),
            T is the telescope/observer position.
        '''
        F = self.FOV
        T = self.T
        return xarray_unit_vector(F - T)

    @format_docstring(**_paramdocs_ray_draping, sub_ntab=1)
    def wsph(self, *, in_box=True):
        '''return wsph, i.e. the points along each ray in sph system cartesian coords centered at r=0.
        
        wsph has dimensions of self.FOV plus 'ray_s' or 'ray_s_dim' dimension,
            as well as 'component' dimension associated with x, y, z coords of each point.

        wsph = F + rayhat * s,   where:
            rayhat = (F - T)/|F - T|,
            F is the point in FOV corresponding to this ray (there is one ray per FOV point),
            T is the telescope/observer position.

        in_box: {in_box}
        '''
        s = self.ray_s_array()
        F = self.FOV
        rayhat = self.rayhat()
        wsph = F + rayhat * s
        if in_box:
            wsph = self.keep_in_box(wsph, in_box=in_box)
        return wsph

    @format_docstring(**_paramdocs_ray_draping, sub_ntab=1)
    def wbox(self, *, in_box=True, periodic=True):
        '''return wbox, i.e. the points along each ray, in box system cartesian coords.
        These are the curved paths which are the "final answer" for the ray draping problem.

        wbox has dimensions of self.FOV plus 'ray_s' or 'ray_s_dim' dimension,
            as well as 'component' dimension associated with x, y, z coords of each point.

        in_box: {in_box}
        periodic: {periodic}
        '''
        wsph = self.wsph(in_box=in_box)
        wbox = self.sph_to_box(wsph, periodic=periodic,
                               in_box=False)  # [EFF] no need to re-check in_box here.
        return wbox

    @format_docstring(**_paramdocs_ray_draping, sub_ntab=1)
    def rays(self, system='box', *, in_box=True, periodic=True):
        '''return rays, i.e. points along each ray, in desired coordinate system.
        Result is an xr.DataArray with dims from self.FOV, 'ray_s' or 'ray_s_dim',
            and 'component' dim associated with x, y, z coords of each point.
    
        system: 'box' or 'sph'
            'box' --> return points in box system cartesian coords.
            'sph' --> return points in sph system cartesian coords centered at r=0.
        in_box: {in_box}
        periodic: {periodic}
            (only used if system=='box')

        Equivalent to self.wbox() or self.wsph(), depending on system.
        '''
        if system == 'box':
            rays = self.wbox(in_box=in_box, periodic=periodic)
        elif system == 'sph':
            rays = self.wsph(in_box=in_box)
        else:
            raise InputError(f"Expected system = 'box' or 'sph'; got {system!r}.")
        return rays

    def ray_tangents(self, system='box', *, unit_vector=False):
        '''return array of vectors pointing in ray directions (i.e. tangent to ray path).
        The implementation here takes derivatives along the 'ray_s' or 'ray_s_dim' dimension,
            which may produce artifacts at edges of ray_s.

        system: 'box' or 'sph'
            'box' --> return directions along rays in box system
            'sph' --> equivalent to self.rayhat() (which is a constant for each ray)
        unit_vector: bool
            whether to convert to unit vector.
            ('sph' result will always be unit vector, regardless.)
            False --> 'box' result will just be derivative of ray with respect to ray_s.
        '''
        if system == 'sph':
            return self.rayhat()
        else:
            rays = self.rays(system=system, periodic=False)
            return xarray_differentiate(rays, 'ray_s')


    # # # "STANDARD" FOV CHOICES # # #

    def _get_box_corners(self, system='box'):
        '''return all 8 corners of the box, as xr.DataArray with 'corner' and 'component' dims.
        Corners are in box system unless system=='sph', in which case converted via self.box_to_sph().
        '''
        BASIS = self._get_basis()
        XHAT = BASIS.XHAT
        YHAT = BASIS.YHAT
        ZHAT = BASIS.ZHAT
        corners = []
        for ix in [0, self.XEXT]:
            for iy in [0, self.YEXT]:
                for iz in [0, self.ZEXT]:
                    corner = ix * XHAT + iy * YHAT + iz * ZHAT
                    corners.append(corner)
        corners = xr.concat(corners, dim='corner')
        if system == 'sph':
            corners = self.box_to_sph(corners)
        return corners

    def _get_box_top(self, system='box', *, Nu=21, Nv=21):
        '''return top surface of box, with Nu x Nv points,
        along dims _pc_surface_u_ and _pc_surface_v_.
        Points are in box system unless system=='sph', in which case converted via self.box_to_sph().
        '''
        uu = xr1d(np.linspace(0, self.XEXT, Nu), name='_pc_surface_u_')
        vv = xr1d(np.linspace(0, self.YEXT, Nv), name='_pc_surface_v_')
        BASIS = self._get_basis()
        result = uu * BASIS.XHAT + vv * BASIS.YHAT + self.ZEXT * BASIS.ZHAT
        if system == 'sph':
            result = self.box_to_sph(result)
        return result

    def standard_fov_limits(self):
        '''return (xmin, xmax, ymin, ymax) for standard FOV.
        This is chosen from scene and anchor params,
            as rectangle (at zsph=R0+ZEXT) with size chosen to cover entire box,
            and edges parallel to x sph and y sph axes.
            It assumes that T is along sph z axis (else crash with NotImplementedError)

        See help(PlasmaCalcs.ray_draping) FOV POSITION AND SIZE for more details.
        See also: self.standard_FOV(), self.standard_fov()
        '''
        # check T along z sph axis.
        T = self.T
        T_hat = xarray_unit_vector(T)
        T_hatx = xarray_sel(T_hat, component='x')
        T_haty = xarray_sel(T_hat, component='y')
        if not np.allclose(T_hatx, 0) or not np.allclose(T_haty, 0):
            raise NotImplementedError("standard_fov_limits() only implemented for T along sph z axis.")
        # projecting to "perp to z" is easy, just ignore the z component.
        # corners
        corners = self._get_box_corners(system='sph')
        corners_x = xarray_sel(corners, component='x')
        x_min = corners_x.min().item()
        x_max = corners_x.max().item()
        corners_y = xarray_sel(corners, component='y')
        y_min = corners_y.min().item()
        y_max = corners_y.max().item()
        # top
        top = self._get_box_top(system='sph')
        top_x = xarray_sel(top, component='x')
        x_min = min(x_min, top_x.min().item())
        x_max = max(x_max, top_x.max().item())
        top_y = xarray_sel(top, component='y')
        y_min = min(y_min, top_y.min().item())
        y_max = max(y_max, top_y.max().item())
        return (x_min, x_max, y_min, y_max)

    def standard_FOV(self, Nx, Ny, *, xdim='FOVx', ydim='FOVy'):
        '''return standard FOV (array) with (Nx, Ny) points.
        This is chosen from scene and anchor params,
            as rectangle (at zsph=R0+ZEXT) with size chosen to cover entire box,
            and edges parallel to x sph and y sph axes.
            It assumes that T is along sph z axis (else crash with NotImplementedError)

        See help(PlasmaCalcs.ray_draping) FOV POSITION AND SIZE for more details.
        See also: self.standard_fov(), self.standard_fov_limits()
        '''
        xmin, xmax, ymin, ymax = self.standard_fov_limits()
        Fx = xr1d(np.linspace(xmin, xmax, Nx), name=xdim)
        Fy = xr1d(np.linspace(ymin, ymax, Ny), name=ydim)
        Fz = self.R0 + self.ZEXT
        BASIS = self._get_basis()
        FOV = Fx * BASIS.XHAT + Fy * BASIS.YHAT + Fz * BASIS.ZHAT
        return FOV

    def standard_fov(self, Nx, Ny, *, xdim='FOVx', ydim='FOVy'):
        '''return standard fov (RayDrapeFOV) with (Nx, Ny) points.
        This is chosen from scene and anchor params,
            as rectangle (at zsph=R0+ZEXT) with size chosen to cover entire box,
            and edges parallel to x sph and y sph axes.
            It assumes that T is along sph z axis (else crash with NotImplementedError)

        See help(PlasmaCalcs.ray_draping) FOV POSITION AND SIZE for more details.
        See also: self.standard_FOV(), self.standard_fov_limits()
        '''
        FOV = self.standard_FOV(Nx, Ny, xdim=xdim, ydim=ydim)
        return self.fov_cls(FOV)

    def _get_fov_value(self, value):
        '''convert value to RayDrapeFOV (unless None) appropriate for storage in self.fov.'''
        if isinstance(value, tuple) and len(value) == 3 and value[0] == 'standard':
            Nx = value[1]
            Ny = value[2]
            value = self.standard_fov(Nx, Ny)
        else:
            value = super()._get_fov_value(value)
        return value


    # # # "STANDARD" RAY CHOICES # # #

    def standard_ray_s_limits(self):
        '''return (min, max) s to use for self.ray_s, corresponding to full overlap with the box.
        Each ray may have its own (min, max) pair, so each is an xarray with same dims as self.FOV.

        See help(PlasmaCalcs.ray_draping) RAY BOUNDARIES section for more details.

        Value will be NaN for rays which never enter the draped box region.
        '''
        rayhat = self.rayhat()
        F = self.FOV
        F_dot_rayhat = xarray_dot_product(F, rayhat)
        F_dot_rayhat_2 = F_dot_rayhat**2
        F2 = xarray_magnitude(F)**2
        # upper boundary:
        r_upper = self.R0 + self.ZEXT
        D_upper = F_dot_rayhat_2 - F2 + r_upper**2  # discriminant
        upper_exists = (D_upper > 0)  # ignore edge case of D==0 for now
        sqrt_D_upper = np.sqrt(D_upper.where(upper_exists, other=np.nan))
        s_upper_entry = - F_dot_rayhat - sqrt_D_upper   # smaller s is always entry
        s_upper_exit  = - F_dot_rayhat + sqrt_D_upper   # larger s is always exit
        # bottom boundary:
        r_lower = self.R0
        D_lower = F_dot_rayhat_2 - F2 + r_lower**2  # discriminant
        lower_exists = (D_lower > 0)  # ignore edge case of D==0 for now
        sqrt_D_lower = np.sqrt(D_lower.where(D_lower >= 0, other=np.nan))
        s_lower_entry = - F_dot_rayhat - sqrt_D_lower
        # s_lower_exit  = - F_dot_rayhat + sqrt_D_lower  # <-- not actually used anywhere
        # combine to get final (min, max).
        s_min = s_upper_entry   # always start with upper entry... no other way to enter.
        s_max = xr.where(
            lower_exists,
            s_lower_entry,   # if lower boundary exists, ray stops at lower entry point.
            s_upper_exit,   # else, ray goes all the way to upper exit point.
        )
        return (s_min, s_max)

    _RAY_S_GRID_DEFAULTS = {
        'name': 'ray_s',
        'inclusive': (False, False),  # avoid including edges, which can lead to NaN during interp.
    }

    def standard_ray_s(self, N):
        '''return standard ray_s with N points along each ray.
        Equivalent: self.ray_s_array() when self.ray_s == N.
        Equivalent: xarray_grid(*self.standard_ray_s_limits(), N=N, **self._RAY_S_GRID_DEFAULTS)
            (self._RAY_S_GRID_DEFAULTS is by default equivalent to: name='ray_s', inclusive=(False, False))
        '''
        smin, smax = self.standard_ray_s_limits()
        return xarray_grid(smin, smax, N=N, **self._RAY_S_GRID_DEFAULTS)

    def ray_s_array(self):
        '''self.ray_s as an xr.DataArray including 'ray_s' or 'ray_s_dim' dimension.
        Behavior depends on self.ray_s:
            xr.DataArray --> if has 'ray_s' or 'ray_s_dim' dim, return as-is, else crash (DimensionalityError)
            None --> crash (InputMissingError)
            int --> use this many points along each ray, with s_min, s_max = self.standard_ray_s_limits()
            dict --> pass kwargs to pc.xarray_grid(), with min, max = self.standard_ray_s_limits() unless provided.
        '''
        result = self.ray_s
        if result is None:
            raise InputMissingError("Cannot get ray_s_array when ray_s is unspecified (self.ray_s=None).")
        elif is_integer(result):
            return self.standard_ray_s(N=result)
        elif isinstance(result, dict):
            if 'min' in result and 'max' in result:
                kw = result
            else:
                smin, smax = self.standard_ray_s_limits()
                kw = {'min': smin, 'max': smax, **self._RAY_S_GRID_DEFAULTS, **result}
            result = xarray_grid(**kw)
        elif not hasattr(result, 'dims'):  # not None, int, dict, nor xarray. Not sure what it is.
            raise InputError("ray_s must be an xarray object including 'ray_s' or 'ray_s_dim' dimension.")
        elif 'ray_s' not in result.dims and 'ray_s_dim' not in result.dims:
            raise DimensionalityError(f"ray_s must have 'ray_s' or 'ray_s_dim' dim, got dims {result.dims}")
        return result

    @property
    def ray_s_dim(self):
        '''return name of ray_s dimension, i.e. 'ray_s' or 'ray_s_dim', depending on which one self.ray_s_array() has.
        [EFF] for now this just checks dims of self.ray_s_array() (i.e. it recomputes that array every time),
            rather than guessing properly before the array is computed.
        '''
        ray_s_array = self.ray_s_array()
        if 'ray_s' in ray_s_array.dims:
            return 'ray_s'
        elif 'ray_s_dim' in ray_s_array.dims:
            return 'ray_s_dim'
        else:
            raise DimensionalityError(f"ray_s_array must have 'ray_s' or 'ray_s_dim' dim, got dims {ray_s_array.dims}")


    # # # PLOTTING # # #

    PLOT_DEFAULTS = {
        ## values related to self.plot() ##
        'box': {
            'bottom': {
                'color': 'orangered',
                'alpha': 0.5,
            },
            'top': {
                'color': 'gray',
                'alpha': 0.5,
            },
            'posts': {
                'color': 'black',
            },
            'A': {
                'color': 'red',
                'marker': r'$\text{A}$',
            },
            'o': {   # marker at xbox=X0, ybox=Y0, zbox=0
                'color': 'orangered',
                'marker': 'o',
                'facecolors': 'none',
            },
            'x': {   # marker at xbox=XEXT, ybox=Y0, zbox=0
                'color': 'orangered',
                'marker': 'x',
            },
        },
        'rays_cmap': 'viridis',
        'rays': {
            'colors': plt.get_cmap('tab10').colors,
        },
        'fov': {
            'color': 'black',
        },
        'view_init': {
            'roll': 270,
            'elev': 15,
            'azim':-15,
        },
        'aspect': (1,1,1),

        ## values related to getting box surfaces/edges/points
        'post_points': 2,  # number of points for each post line connecting box corners
        'surface_points': (21, 21),  # number of points in (u, v) grid for top & bottom
    }

    @format_docstring(**_paramdocs_ray_draping)
    def plot(self, system='both', *, box=True, fov=True, rays=True, max_n_rays=100,
             in_box='lower', periodic=True,
             ax=None, axsize=(6,6), aspect=True, view_init=True,
             extra_points=None, extra_lines=None):
        '''visualize self by plotting. Useful for sanity checks!
        Returns matplotlib ax (or tuple of (ax_sph, ax_box) if system='both').

        If using Jupyter, for easier 3D visualization consider running a cell with command:
            %matplotlib notebook 

        (Below, False-ness checked via "is False", to avoid ambiguity with empty dicts)

        system: 'both', 'sph' or 'box'
            which system(s) to visualize.
            'both' --> make a figure with two plots, one for each system.
            'sph' --> 'sph' system, where rays are straight, and box top and bottom are curved.
            'box' --> 'box' system, where rays are curved, and box top and bottom are flat.
        box: bool or dict of dicts
            whether to plot the box boundaries.
            True --> use default box plotting parameters (self.PLOT_DEFAULTS['box']).
            dict of dicts --> customize box plotting parameters. vals are dict or False.
                Use False to indicate "don't plot this part".
                Skip key to indicate "use default kwargs for this part".
                'bottom': kwargs for ax.plot_surface() for bottom of box (zbox=0 plane)
                'top': kwargs for ax.plot_surface() for top of box (zbox=ZEXT plane)
                'posts': kwargs for ax.plot() for vertical posts connecting bottom/top corners
                'A': kwargs for ax.scatter() for anchor point
                'o': kwargs for ax.scatter() for point at xbox=X0, ybox=Y0, zbox=0
                'x': kwargs for ax.scatter() for point at xbox=XEXT, ybox=Y0, zbox=0
        fov: bool or dict
            whether to plot the FOV points.
            True --> use default FOV plotting parameters (self.PLOT_DEFAULTS['fov']).
            dict --> customize FOV plotting parameters. Can provide any kwargs for ax.scatter().
        rays: bool or dict
            whether to plot the rays.
            True --> use default ray plotting parameters (self.PLOT_DEFAULTS['box']).
            dict --> customize ray plotting parameters. Can provide any kwargs for ax.plot().
            Some kwargs are treated specially: 'colors' and 'cmap'
                if 'colors' is provided, cycle these colors if there are more colors than rays,
                else use 'cmap', or self.PLOT_DEFAULTS['rays_cmap'] if 'cmap' not provided.
        max_n_rays: int or None
            maximum number of rays to plot, else crash (if doing fov and/or rays).
            prevents accidentally plotting way too much if doing production-grade RayDraping.
            None --> no limit.
        in_box: {in_box}
            (The check here applies only to plotted rays)
        periodic: {periodic}
            (The wrapping here applies only to plotted rays and FOV points when system=='box')
            (If periodic wrapping leads to weird-looking lines crossing the box,
                consider using rays={{'linestyle':'none', 'marker':'.'}} to just plot points.)
        ax: None or matplotlib Axes3D object
            Axes to plot on. None --> create new figure and axes.
            Incompatible with system='both'.
        axsize: tuple of two ints
            (x,y) size of new figure, per axes (only if making new axes because ax is None).
            figsize = (2x, y) if system='both', else (x,y).
        aspect: bool or 3-tuple of numbers
            whether to set aspect ratio for 3D axes.
            True --> set aspect based on self.PLOT_DEFAULTS['view_init']. (default: (1,1,1))
            3-tuple --> (x,y,z) scaling to use.
            (Matplotlib is annoying about this, so we have to do:
                sizes=[xmax-xmin, ymax-ymin, zmax-zmin] from get_xlim(), get_ylim(), get_zlim(),
                ax.set_box_aspect(sizes[0]*aspect[0], sizes[1]*aspect[1], sizes[2]*aspect[2]))
        view_init: bool or dict
            whether to ax.view_init().
            True --> use self.PLOT_DEFAULTS['view_init'].
            dict --> dict of kwargs for ax.view_init(). (Consider: 'elev', 'azim', 'roll'.)
        '''
        # checks first (if will crash, crash before plotting anything)
        if system not in ['both', 'sph', 'box']:
            raise InputError(f"system must be 'both', 'sph', or 'box', but got {system!r}")
        if ax is not None and system == 'both':
            raise InputConflictError("Cannot provide ax when system='both'.")
        if (max_n_rays is not None) and (fov is not False or rays is not False):
            if self.n_rays > max_n_rays:
                errmsg = (f'Too many rays to plot: {self.n_rays}. Max is max_n_rays={max_n_rays}.'
                          '\nConsider changing max_n_rays, or reducing number of points in FOV.'
                          '\nUse max_n_rays=None to remove this limit entirely.')
                raise TooManyPlottablesError(errmsg)
        # setup axes
        if ax is None:
            if system == 'both':
                fig = plt.figure(figsize=(axsize[0]*2, axsize[1]))
                ax_sph = fig.add_subplot(1, 2, 1, projection='3d')
                ax_box = fig.add_subplot(1, 2, 2, projection='3d')
                kw_plot = dict(box=box, fov=fov, rays=rays, max_n_rays=max_n_rays,
                               in_box=in_box, periodic=periodic,
                               aspect=aspect, view_init=view_init)
                ax_sph = self.plot('sph', ax=ax_sph, **kw_plot)
                ax_box = self.plot('box', ax=ax_box, **kw_plot)
                return (ax_sph, ax_box)
            else:
                fig = plt.figure(figsize=axsize)
                ax = fig.add_subplot(1, 1, 1, projection='3d')
        assert system != 'both'    # if this assertion fails, made a coding error.
        # do plotting
        if box is not False:
            boxdefaults = self.PLOT_DEFAULTS['box']
            boxkeys = ['bottom', 'top', 'posts', 'o', 'x']
            if isinstance(box, dict):
                boxkws = {k: {**boxdefaults[k], **box.get(k, {})} for k in boxkeys}
            else:
                boxkws = boxdefaults
            boxkws = {k: v for k, v in boxkws.items() if v is not False}
            boxv = self._get_box_plot_values(system=system, include=boxkws.keys())
            if 'A' in boxkws:
                ax.scatter(*xarray_take_xyz(boxv['A']), **boxkws['A'])
            if 'o' in boxkws:
                ax.scatter(*xarray_take_xyz(boxv['o']), **boxkws['o'])
            if 'x' in boxkws:
                ax.scatter(*xarray_take_xyz(boxv['x']), **boxkws['x'])
            if 'posts' in boxkws:
                for post_key in ['post00', 'post10', 'post01', 'post11']:
                    ax.plot(*xarray_take_xyz(boxv[post_key]), **boxkws['posts'])
            if 'bottom' in boxkws:
                ax.plot_surface(*xarray_take_xyz(boxv['bottom']), **boxkws['bottom'])
            if 'top' in boxkws:
                ax.plot_surface(*xarray_take_xyz(boxv['top']), **boxkws['top'])
        if fov is not False:
            fovdefaults = self.PLOT_DEFAULTS['fov']
            fovkws = {**fovdefaults, **fov} if isinstance(fov, dict) else fovdefaults
            FOVv = self.FOV if system == 'sph' else self.sph_to_box(self.FOV, periodic=periodic)
            ax.scatter(*xarray_take_xyz(FOVv), **fovkws)
        if rays is not False:
            raydefaults = self.PLOT_DEFAULTS['rays']
            raykws = {**raydefaults, **rays} if isinstance(rays, dict) else {**raydefaults}
            colors = raykws.pop('colors', None)
            cmap = raykws.pop('cmap', self.PLOT_DEFAULTS['rays_cmap'])
            n_rays = self.n_rays
            if colors is not None:
                if len(colors) < n_rays:
                    colors = None  # not enough colors
            if colors is None:
                cmap = plt.get_cmap(cmap)
                colors = cmap(np.linspace(0, 1, n_rays))
            if system == 'sph':
                rayv = self.wsph(in_box=in_box)
            else:
                rayv = self.wbox(in_box=in_box, periodic=periodic)
            raystack = rayv.stack(_pc_ray_i_=list(set(rayv.dims) - {'component', 'ray_s', 'ray_s_dim'}))
            for i in range(n_rays):
                ray = raystack.isel(_pc_ray_i_=i)
                ax.plot(*xarray_take_xyz(ray), color=colors[i], **raykws)
        # view-related stuff
        ax.set(xlabel='x', ylabel='y', zlabel='z', title=f'{system} system')
        if aspect is not False:
            if aspect is True:
                aspect = self.PLOT_DEFAULTS['aspect']
            lims = [getattr(ax, f'get_{x}lim')() for x in ['x', 'y', 'z']]
            sizes = [lim[1] - lim[0] for lim in lims]
            ax.set_box_aspect((sizes[0]*aspect[0], sizes[1]*aspect[1], sizes[2]*aspect[2]))
        if view_init is not False:
            if not isinstance(view_init, dict):
                view_init = self.PLOT_DEFAULTS['view_init']
            ax.view_init(**view_init)
        return ax

    def _get_box_plot_values(self, *, system='box', include=['o', 'x', 'A', 'posts', 'bottom', 'top']):
        '''return dict of values related to box, for self.plot(box=True).
        Results are xr.DataArrays with 'component' dim telling x,y,z coords.

        system: 'box' or 'sph'
            which system to get box values for.
            'box' --> returns values in box system cartesian coords. No conversion necessary.
            'sph' --> convert to sph system cartesian coords centered at r=0, via self.box_to_sph().
        include: list of strs, tells which keys to include. (Use 'posts' to include all posts):
            'o': point at xbox=X0, ybox=Y0, zbox=0. This is the lower-left point in the box.
            'x': point at xbox=X0+XEXT, ybox=Y0, zbox=0
            'A': the anchor point
            'post00': (xbox=X0,ybox=Y0), zbox from 0 to ZEXT, with dim '_pc_post_i_'
            'post10': (xbox=X0+XEXT,ybox=Y0), zbox from 0 to ZEXT, with dim '_pc_post_i_'
            'post01': (xbox=X0,ybox=Y0+YEXT), zbox from 0 to ZEXT, with dim '_pc_post_i_'
            'post11': (xbox=X0+XEXT,ybox=Y0+YEXT), zbox from 0 to ZEXT, with dim '_pc_post_i_'
            'bottom': zbox=0 surface, with dims '_pc_surface_u_', and '_pc_surface_v_'
            'top': zbox=ZEXT surface, with dims '_pc_surface_u_', and '_pc_surface_v_'

        The number of points along the posts and surface dims are determined by
            self.PLOT_DEFAULTS['post_points'] and self.PLOT_DEFAULTS['surface_points'].
        '''
        result = {}
        BASIS = self._get_basis()
        XHAT = BASIS.XHAT
        YHAT = BASIS.YHAT
        ZHAT = BASIS.ZHAT
        X0 = self.X0
        Y0 = self.Y0
        XEXT = self.XEXT
        YEXT = self.YEXT
        X1 = X0 + XEXT
        Y1 = Y0 + YEXT
        # points at (0,0,0) and (XEXT,0,0)
        if 'o' in include:
            result['o'] = X0 * XHAT + Y0 * YHAT
        if 'x' in include:
            result['x'] = X1 * XHAT + Y0 * YHAT
        if 'A' in include:
            result['A'] = self.Abox_x * XHAT + self.Abox_y * YHAT   # + 0 * ZHAT
        # posts
        if 'posts' in include:
            pspanz = np.linspace(0, self.ZEXT, self.PLOT_DEFAULTS['post_points'])
            pspanz = xr1d(pspanz, name='_pc_post_i_')
            xyposts = {'post00': (X0,Y0), 'post10': (X1,Y0),
                       'post01': (X0,Y1), 'post11': (X1,Y1)}
            result.update({k: x_ * XHAT + y_ * YHAT + pspanz * ZHAT for k, (x_, y_) in xyposts.items()})
        # bottom and top surfaces
        if 'bottom' in include or 'top' in include:
            uu = np.linspace(X0, X1, self.PLOT_DEFAULTS['surface_points'][0])
            vv = np.linspace(Y0, Y1, self.PLOT_DEFAULTS['surface_points'][1])
            uu = xr1d(uu, name='_pc_surface_u_')
            vv = xr1d(vv, name='_pc_surface_v_')
            if 'bottom' in include:
                result['bottom'] = uu * XHAT + vv * YHAT + 0 * ZHAT
            if 'top' in include:
                result['top'] = uu * XHAT + vv * YHAT + self.ZEXT * ZHAT
        # convert system if needed
        if system == 'sph':
            for k, v in result.items():
                result[k] = self.box_to_sph(v)
        return result

    def get_post(self, x, y, system='box', *, dim='_pc_post_i_', N=2):
        '''returns this post, i.e. points corresponding to line from (x,y,0) to (x,y,ZEXT) in box system.
        Result is an xarray.DataArray with 'component' dimension, converted to the indicated system.

        x, y: scalars
            coordinates of post in box system.
            In box system, the post is a line from (x,y,0) to (x,y,ZEXT).
        dim: str
            name of dimension for points along post.
        N: int
            number of points along post (including endpoints).
        '''
        BASIS = self._get_basis()
        zz = xr1d(np.linspace(0, self.ZEXT, N), name=dim)
        post = x * BASIS.XHAT + y * BASIS.YHAT + zz * BASIS.ZHAT
        if system == 'sph':
            post = self.box_to_sph(post)
        return post

    def plot_post(self, x, y, ax_sph=None, ax_box=None, **kw_plot):
        '''plot a post (i.e. line from z=0 to z=ZEXT in box system) on the indicated ax(es).

        x, y: scalars
            coordinates of post in box system.
            In box system, the post is a line from (x,y,0) to (x,y,ZEXT).
        ax_sph, ax_box: None or matplotlib axes objects
            plots post in sph system on ax_sph, in box system on ax_box.
            None --> don't plot on that system.
            At least one of these two inputs must be provided, else crash.
        Additional kwargs go to matplotlib ax.plot() calls.

        Example (plots post at x=5, y=7 in both systems):
            draper = RayDraping(...)
            axes = draper.plot(system='both')
            draper.plot_post(5, 7, *axes)
            # equivalent: draper.plot_post(5, 7, ax_sph=axes[0], ax_box=axes[1])
        '''
        if ax_sph is None and ax_box is None:
            raise InputMissingError("Must provide at least one of ax_sph or ax_box.")
        if ax_sph is not None:
            post_sph = self.get_post(x, y, system='sph')
            ax_sph.plot(*xarray_take_xyz(post_sph), **kw_plot)
        if ax_box is not None:
            post_box = self.get_post(x, y, system='box')
            ax_box.plot(*xarray_take_xyz(post_box), **kw_plot)


    # # # DISPLAY # # #

    def _repr_contents(self):
        '''list of contents for self.__repr__'''
        result = super()._repr_contents()
        if is_integer(self.ray_s) or isinstance(self.ray_s, dict):
            result.append(f"ray_s={self.ray_s}")
        elif self.ray_s is not None:
            result.append(f"ray_s=DataArray(sizes={dict(self.ray_s.sizes)})")
        return result
