LON_STR = 'lon' LAT_STR = 'lat' PFULL_STR = 'pfull' class PhysDeriv(object):
    """Derivatives in physical space."""
    _COORD_CLS = Coord
    _DERIV_CLS = FiniteDeriv
    _WRAP_CIRCUMF = 0
    _WRAP_LEFT_TO_RIGHT = 0
    _WRAP_RIGHT_TO_LEFT = 0

    def _get_coord(self, coord):
        if coord is None:
            return self.arr[self.dim].copy(deep=True)
        return coord

    def __init__(self, arr, dim, coord=None, spacing=1, order=2,
                 fill_edge=True, **coord_kwargs):
        self.arr = arr.copy(deep=True)
        self.dim = dim
        self.coord = self._get_coord(coord)
        self._orig_coord_values = self.coord.values
        self.spacing = spacing
        self.order = order
        self.cyclic = coord_kwargs.get('cyclic', False)

        self._coord_obj = self._COORD_CLS(self.coord, dim=self.dim,
                                          **coord_kwargs)
        for attr in ['deriv_factor', 'deriv_prefactor']:
            setattr(self, attr, getattr(self._coord_obj, attr))

        if fill_edge and self.cyclic:
            warnings.warn("Overriding 'fill_edge' value of True, because "
                          "coord is cyclic")
        self.fill_edge = False if self.cyclic else fill_edge

    def _wrap(self, arr):
        if self.cyclic:
            return wraparound(
                arr.copy(deep=True), self.dim,
                left_to_right=self._WRAP_LEFT_TO_RIGHT*self.order,
                right_to_left=self._WRAP_RIGHT_TO_LEFT*self.order,
                circumf=self._WRAP_CIRCUMF, spacing=self.spacing
            )
        return arr.copy(deep=True)

    def _prep_coord(self, coord):
        return coord

    def deriv(self, *args, **kwargs):
        """Derivative, incorporating physical/geometrical factors."""
        arr = self._wrap(self.arr.copy(deep=True) *
                         self.deriv_factor(*args, **kwargs))
        coord = self._prep_coord(arr[self.dim].copy(deep=True))
        darr = (self._DERIV_CLS(arr.copy(deep=True), self.dim,
                                coord=coord.copy(deep=True),
                                spacing=self.spacing, order=self.order,
                                fill_edge=self.fill_edge).deriv() *
                self._coord_obj.deriv_prefactor(*args, **kwargs))
        return darr class LonFwdDeriv(LonDeriv):
    _DERIV_CLS = FwdDeriv
    _WRAP_LEFT_TO_RIGHT = 1 class LonBwdDeriv(LonDeriv):
    _DERIV_CLS = BwdDeriv
    _WRAP_RIGHT_TO_LEFT = 1 class LatFwdDeriv(LatDeriv):
    _DERIV_CLS = FwdDeriv class LatBwdDeriv(LatDeriv):
    _DERIV_CLS = BwdDeriv class EtaFwdDeriv(EtaDeriv):
    _DERIV_CLS = FwdDeriv class EtaBwdDeriv(EtaDeriv):
    _DERIV_CLS = BwdDeriv class SphereEtaBwdDeriv(SphereEtaDeriv):
    """Derivatives on the sphere with hybrid sigma-pressure in the vertical."""
    _HORIZ_DERIV_CLS = SphereBwdDeriv
    _VERT_DERIV_CLS = EtaBwdDeriv class SphereEtaFwdDeriv(SphereEtaDeriv):
    """Derivatives on the sphere with hybrid sigma-pressure in the vertical."""
    _HORIZ_DERIV_CLS = SphereFwdDeriv
    _VERT_DERIV_CLS = EtaFwdDeriv class Upwind(Advec):
    """Upwind advection."""
    _DERIV_BWD_CLS = BwdDeriv
    _DERIV_FWD_CLS = FwdDeriv
    _DERIV_METHOD = 'deriv'

    def __init__(self, flow, arr, dim, coord=None, spacing=1, order=2,
                 fill_edge=True):
        super(Upwind, self).__init__(flow, arr, dim, coord=coord,
                                     spacing=spacing, order=order,
                                     fill_edge=fill_edge)
        self._deriv_bwd_obj = self._DERIV_BWD_CLS(
            self.arr, self.dim, coord=self.coord, spacing=self.spacing,
            order=self.order, fill_edge=True
        )
        self._deriv_fwd_obj = self._DERIV_FWD_CLS(
            self.arr, self.dim, coord=self.coord, spacing=self.spacing,
            order=self.order, fill_edge=True
        )
        self._deriv_bwd = getattr(self._deriv_bwd_obj, self._DERIV_METHOD)
        self._deriv_fwd = getattr(self._deriv_fwd_obj, self._DERIV_METHOD)

    def _flow_neg_pos(self, reverse_dim=False):
        """Create negative- and positive-only arrays for upwind advection.

        :out: flow_neg, flow_pos xarray.DataArrays with shape and coords
            identical to `flow1, but with, respectively, all positive and
            negative values set to 0 (or the reverse if `reverse_dim` is
            `True`).
        """
        flow_neg = self.flow.copy()
        flow_neg.values[self.flow.values >= 0] = 0.
        flow_pos = self.flow.copy()
        flow_pos.values[self.flow.values < 0] = 0.
        if not reverse_dim:
            return flow_neg, flow_pos
        return flow_pos, flow_neg

    def _swap_bwd_fwd_edges(self, bwd, fwd):
        """Forward diff on left edge; backward diff on right edge."""
        edge_left = {self.dim: 0}
        edge_right = {self.dim: -1}
        bwd[edge_left] = fwd[edge_left]
        fwd[edge_right] = bwd[edge_right]
        return bwd, fwd

    def _derivs_bwd_fwd(self):
        """Generate forward and backward differencing derivs for upwind.

        Order of accuracy decreases moving towards edge (right edge for
        forward, left edge for backward) as the differencing stencil starts to
        extend over the domain edge.  At the edge itself, the opposite signed
        differencing is used with the same order of accuracy as in the
        interior.
        """
        bwd = self._deriv_bwd()
        fwd = self._deriv_fwd()
        # Forward diff on left edge; backward diff on right edge.
        return self._swap_bwd_fwd_edges(bwd, fwd)

    def advec(self):
        """
        Upwind differencing scheme for advection.

        In interior, forward differencing for negative flow, and backward
        differencing for positive flow.

        :param arr: Field being advected.
        :param flow: Flow that is advecting the field.
        """
        bwd, fwd = self._derivs_bwd_fwd()
        neg, pos = self._flow_neg_pos()
        advec_arr = pos*bwd + neg*fwd
        if not self.fill_edge:
            slice_middle = {self.dim: slice(self.order, -self.order)}
            advec_arr = advec_arr[slice_middle]
        return advec_arr
from .._constants import LON_STR, LAT_STR, PFULL_STR
from ..deriv import (PhysDeriv, LonBwdDeriv, LonFwdDeriv, LatBwdDeriv,
                     LatFwdDeriv, EtaBwdDeriv, EtaFwdDeriv,
                     SphereEtaBwdDeriv, SphereEtaFwdDeriv)
from . import Upwind


def _make_derivs(obj, arr, *deriv_args, **deriv_kwargs):
    obj._deriv_bwd_obj = obj._DERIV_BWD_CLS(arr, *deriv_args, **deriv_kwargs)
    obj._deriv_bwd = getattr(obj._deriv_bwd_obj, obj._DERIV_METHOD)
    obj._deriv_fwd_obj = obj._DERIV_FWD_CLS(arr, *deriv_args, **deriv_kwargs)
    obj._deriv_fwd = getattr(obj._deriv_fwd_obj, obj._DERIV_METHOD)


class PhysUpwind(Upwind):
    """Upwind advection along a physical coordinate."""
    _DERIV_BWD_CLS = PhysDeriv
    _DERIV_FWD_CLS = PhysDeriv
    _DERIV_METHOD = 'deriv'
    _ADVEC_CLS = Upwind

    def __init__(self, flow, arr, dim, coord=None, spacing=1, order=2,
                 cyclic=False, fill_edge=True):
        self.flow = flow
        self.arr = arr
        self.dim = dim
        self.coord = coord
        self.spacing = spacing
        self.order = order
        self.cyclic = cyclic
        self.fill_edge = fill_edge

        deriv_args = [dim]
        deriv_kwargs = dict(coord=coord, spacing=spacing, order=order,
                            cyclic=cyclic, fill_edge=fill_edge)
        _make_derivs(self, arr, *deriv_args, **deriv_kwargs)

    def _derivs_bwd_fwd(self, *args, **kwargs):
        """Generate forward and backward differencing derivs for upwind.

        If not cyclic, order of accuracy decreases moving towards edge (right
        edge for forward, left edge for backward) as the differencing stencil
        starts to extend over the domain edge.  At the edge itself, the
        opposite signed differencing is used with the same order of accuracy as
        in the interior.
        """
        bwd = self._deriv_bwd(*args, **kwargs)
        fwd = self._deriv_fwd(*args, **kwargs)
        # Forward diff on left edge; backward diff on right edge if not cyclic.
        if getattr(self, 'cyclic', False):
            return bwd, fwd
        return self._swap_bwd_fwd_edges(bwd, fwd)

    def advec(self, *args, **kwargs):
        """
        Upwind differencing scheme for advection.

        In interior, forward differencing for negative flow, and backward
        differencing for positive flow.

        :param arr: Field being advected.
        :param flow: Flow that is advecting the field.
        """
        bwd, fwd = self._derivs_bwd_fwd(*args, **kwargs)
        neg, pos = self._flow_neg_pos()
        advec_arr = pos*bwd + neg*fwd
        if not self.fill_edge and not self.cyclic:
            slice_middle = {self.dim: slice(self.order, -self.order)}
            advec_arr = advec_arr[slice_middle]
        return advec_arr


class LonUpwind(PhysUpwind):
    """Upwind advection in longitude."""
    _DERIV_BWD_CLS = LonBwdDeriv

    _DERIV_FWD_CLS = LonFwdDeriv