"""
Methods for the probability integral transform (PIT) classes Pit and PitFcstAtObs.

The implementation follows the theoretical approach of Gneiting and Ranjan (2013) and
Taggart (2023) (see https://nla.gov.au/nla.obj-3079961862/view) which interprets the PIT
of a forecast--observation pair as a CDF. Details of this approach are described in this docstring.

For a forecast--observation pair (G,y), where G is a CDF,
the corresponding PIT value is the uniform distribution on the closed interval
[G(y-), G(y)], where G(y-) denotes the left-hand limit.
So if F(G-)=G(y), the PIT value is the distribution concentrated at the point G(y).
PIT values for particular forecast cases can be represented in the form
    lower = G(y-), upper = G(y)
or can be described by the CDF of the uniform distribution on the interval [G(y-), G(y)].

The code here first converts inputs to the [lower, upper] representation.
Then it converts this representation to the uniform CDF representation, with CDF values
stored at a minimal number of points. The remaining values can be obtained exactly when
required via linear interpolation.

Once each forecast case has a PIT value in CDF representation, the (possibly weighted)
mean PIT value across all forecast cases is the weighted mean across all the CDF representations.
We call this the PIT distribution (or PIT CDF) for all the forecast cases.
This PIT CDF is given the symbol F in docstrings throughout this module.

The PIT distribution F is always piecewise linear and right-continuous. It can be represented
using "left" and "right" values, respectively representing the left-hand limit of the PIT CDF
and the value (which equals the right-hand limit) of the CDF. All standard statistics of
PIT for the set of forecast cases, such as PIT histogram bar heights and alpha scores,
can be calculated exactly from this representation.

The code here is structured as follows:
1. The two classes `Pit` and `PitFcstAtObs` are introduced
2. Private functions that calculate the PIT distribution F are then given
3. Private functions that calculate the PIT statistics from the PIT distribution F are then presented.

The tutorial is a good place to obtain an overview of the big picture.
"""

import warnings
from typing import Hashable, Optional

import numpy as np
import xarray as xr
from scipy.interpolate import interp1d

from scores.probability.checks import cdf_values_within_bounds, coords_increasing
from scores.processing import aggregate, broadcast_and_match_nan
from scores.processing.cdf import add_thresholds
from scores.typing import FlexibleDimensionTypes, XarrayLike
from scores.utils import gather_dimensions

RESERVED_NAMES = {
    "uniform_endpoint",
    "pit_x_value",
    "x_plotting_position",
    "y_plotting_position",
    "plotting_point",
    "bin_left_endpoint",
    "bin_right_endpoint",
    "bin_centre",
}

################################################
# The two classes: `Pit`` and `PitFcstAtObs`
################################################


class Pit:
    """
    Calculates the probability integral transform (PIT) for a set of forecast cases and
    corresponding observations. The calculated PIT can be a (possibly weighted) average
    over specified dimensions. Calculations are performed and interpreted as follows.

    Given a predictive cumulative distribution function (CDF) :math:`G` and corresponding
    observation :math:`y`, the corresponding PIT value is a uniform distribution on the
    closed interval :math:`[G(y-), G(y)]`, where :math:`G(y-)` denotes the left-hand
    limit of :math:`G` at :math:`y`. This is the most general form of PIT and handles cases
    where the observation :math:`y` coincides with a point of discontinuity of a predictive CDF
    :math:`G`. See Gneiting and Ranjan (2013) and Taggart (2022).

    Weighted means of PIT values are simply weighted means of CDFs of those uniform distributions
    on the closed interval :math:`[G(y-), G(y)]`. The weighted mean is itself a CDF, here denoted :math:`F`
    and called the PIT distribution for the set of forecast cases.
    Left-hand and right-hand limits of :math:`F` are accessible via the attributes ``.left`` and ``.right``.
    Important statistics related to the PIT distribution :math:`F`, such as bar heights of PIT histograms,
    the graph of :math:`F`, plotting points for PIT uniform probability plots, the alpha score, and
    expected value and variance of :math:`F`, can be calculated from :math:`F` and are accessible via
    methods on the ``Pit`` class.

    Forecasts arguments to the ``Pit`` class can be given in two forms:

    - An ensemble for forecasts, indexed by an ensemble member dimension.
        In this case, the ensemble is interpreted as the empirical cumulative distribution
        function (CDF) generated by the ensemble members.
    - Values of the predictive CDF, indexed by a cdf threshold dimension.
        The values :math:`G(x)` of the predictive CDF :math:`G` evaluated at thresholds :math:`x` must be supplied.
        Left-hand limits :math:`G(x-)` of the predictive CDF can also be specified, allowing for
        CDFs with discontinuities. Values of the CDF between threshold coordinates :math:`x`
        are determined via linear interpolation, whilst values of the CDF at points outside
        those coordinates are assigned either 0 or 1 as appropriate. Any predictive CDF with
        a NaN value will be treated as NaN in its entirety.

    Consider using the ``PitFcstAtObs`` class instead of ``Pit`` when values of the
    predictive CDFs :math:`G` evaluated at the observations are easy to calculate,
    e.g. when the predictive CDFs are normal distributions with known parameters.

    Args:
        fcst: an xarray object of forecasts, either in giving values of the ensemble members or of the predictive CDF.
            If an ensemble, must contain the supplied dimension ``ensemble_member_dim``.
            If a CDF, must contain the supplied dimension ``cdf_threshold_dim``.
        obs: an xarray object of observations.
        ensemble_member_dim: name of the dimension in ``fcst`` that indexes the ensemble members
            if ``fcst`` is an ensemble; otherwise ``None``.
        cdf_threshold_dim: name of the dimension in ``fcst`` that indexes the thresholds
            if ``fcst`` contains predictive CDF values; otherwise ``None``.
        fcst_left: The values of the left-hand limits of the predictive CDF. Must have the same
            shape and dimensions as ``fcst``.
            Only required when ``cdf_threshold_dim`` has been supplied and
            the predictive CDF :math:`G` is discontinuous.
        reduce_dims: Optionally specify which dimensions to reduce when calculating the
            PIT CDF values, where the mean is taken over all forecast cases.
            All other dimensions will be preserved. As a special case, 'all' will allow
            all dimensions to be reduced. Only one of ``reduce_dims`` and ``preserve_dims``
            can be supplied. The default behaviour if neither are supplied is to reduce all dims.
        preserve_dims: Optionally specify which dimensions to preserve when calculating the
            PIT CDF values, where the mean is taken over all forecast cases.
            All other dimensions will be reduced. As a special case, 'all' will allow
            all dimensions to be preserved, apart from ``special_fcst_dim``.
            Only one of ``reduce_dims`` and ``preserve_dims`` can be supplied. The default
            behaviour if neither are supplied is to reduce all dims.
        weights: Optionally provide an array for weighted averaging (e.g. by area, by latitude,
            by population, custom) of PIT CDF values across all forecast cases.

    Attributes:
        left: xarray object giving the left-hand limits :math:`F(x-)` of the CDF :math:`F` of the
            (possibly weighted) PIT distribution across all forecast cases. The points :math:`x`
            are in the dimension "pit_x_value". The CDF :math:`F` of the PIT distribution is completely
            reconstructable using values from ``.left``, ``.right`` and linear interpolation.
        right: xarray object giving the values :math:`F(x)` of the CDF :math:`F` of the
            (possibly weighted) PIT distribution across all forecast cases. The points :math:`x`
            are in the dimension "pit_x_value". The CDF :math:`F` of the PIT distribution is completely
            reconstructable using values from ``.left``, ``.right`` and linear interpolation.
        pit_uniform_endpoints: xarray object giving the interval :math:`[G(y-), G(y)]`
            of the uniform distribution generated from the probability integral transform for each
            forecast--observation pair :math:`(G,y)`. Dimensions are not reduced and weights are ignored.
            The left and right endpoints of the interval are given by the "lower" and "upper" coordinates
            in the "uniform_endpoint" dimension.

    Raises:
        ValueError: when exactly one string is not supplied to either ``ensemble_member_dim`` or ``cdf_threshold_dim``.
        ValueError: when dimensions of ``fcst``, ``obs`` or ``weights`` contain any of the following reserved names:
            "uniform_endpoint", "pit_x_value", "x_plotting_position", "y_plotting_position", "plotting_point",
            "bin_left_endpoint", "bin_right_endpoint", "bin_centre".
        ValueError: when ``fcst_left`` is not ``None``,  and ``fcst_left`` and ``fcst`` do not have identical
            shape, dimensions and coordinates.
        ValueError: when ``cdf_threshold_dim`` has been supplied and any values of
            ``fcst`` that are less then 0 or greater than 1.
        ValueError: when ``cdf_threshold_dim`` has been supplied, ``fcst_left`` is not ``None`` and any
            values of ``fcst_left`` are less then 0 or greater than 1.
        ValueError: when ``cdf_threshold_dim`` has been supplied, ``fcst_left`` is not ``None`` and
            any values of ``fcst_left`` are greater than ``fcst``.
        ValueError: when ``cdf_threshold_dim`` has been supplied and ``fcst[cdf_threshold_dim]`` is not increasing.

    Warns:
        UserWarning: when ``cdf_threshold_dim`` has been supplied, and any values of ``obs`` are less than
            the minimum of ``fcst[cdf_threshold_dim]`` or greater than the maximum of
            ``fcst[cdf_threshold_dim]``.
        UserWarning: when ``cdf_threshold_dim`` has been supplied and there are any NaN values in ``fcst`` or
            ``fcst_left``.

    References:
        - Gneiting, T., & Ranjan, R. (2013). Combining predictive distributions. Electron. J. Statist. 7: 1747-1782 \
            https://doi.org/10.1214/13-EJS823
        - Taggart, R. J. (2022). Assessing calibration when predictive distributions have discontinuities. \
            Bureau Research Report 64, https://nla.gov.au/nla.obj-3079961862/view

    See also:
            - :py:func:`scores.probability.PitFcstAtObs`
            - :py:func:`scores.probability.rank_histogram`

    Examples:
        Calculate the PIT for an under-dispersive ensemble forecast, and calculate various diagnostics and statistics.

        >>> import xarray as xr
        >>> from scipy.stats import norm
        >>> from scores.probability import Pit
        >>> # generate the forecasts and observations then calculate PIT
        >>> fcst = xr.DataArray(norm.rvs(size=(500, 10)), dims=['time', 'ensemble_member'])
        >>> obs = xr.DataArray(norm.rvs(scale=2, size=(500)), dims=['time'])
        >>> pit = Pit(fcst, obs, ensemble_member_dim='ensemble_member')
        >>> # bar heights for a PIT histogram
        >>> histogram_values = pit.hist_values(10)
        >>> # plot the CDF of the PIT distribution
        >>> pit.plotting_points().plot()
        >>> # the expected value of the PIT distribution
        >>> pit_ev = pit.expected_value()
        >>> # the variance of the PIT distribution
        >>> pit_var = pit.variance()
        >>> # the alpha score of the PIT distribution
        >>> pit_alpha = pit.alpha_score()

        Other examples are found in the ``scores`` tutorial for PIT.
    """

    def __init__(
        self,
        fcst: XarrayLike,
        obs: XarrayLike,
        *,  # Force keywords arguments to be keyword-only
        ensemble_member_dim: Optional[str] = None,
        cdf_threshold_dim: Optional[str] = None,
        fcst_left: Optional[XarrayLike] = None,
        reduce_dims: Optional[FlexibleDimensionTypes] = None,
        preserve_dims: Optional[FlexibleDimensionTypes] = None,
        weights: Optional[XarrayLike] = None,
    ):
        """
        Calculates the left-hand and right-hand limits of the PIT distribution :math:`F`,
        accessible via the attributes ``.left`` and ``.right``, as well as the interval
        :math:`[G(y-),G(y)]` of each probability integral transformed forecast--observation pair
        :math:`(G,y)`, accessible via the attribute ``.pit_uniform_endpoints``.
        """
        special_dims = [ensemble_member_dim, cdf_threshold_dim]
        if sum(x is not None for x in special_dims) != 1:
            raise ValueError(
                "A string must be supplied to exactly one of `ensemble_member_dim` and `cdf_threshold_dim`"
            )

        if ensemble_member_dim is not None:
            pit_cdf = _pit_distribution_for_ens(
                fcst,
                obs,
                ensemble_member_dim,
                reduce_dims=reduce_dims,
                preserve_dims=preserve_dims,
                weights=weights,
            )
        else:
            pit_cdf = _pit_distribution_for_cdf(
                fcst,
                obs,
                cdf_threshold_dim,
                fcst_left=fcst_left,
                reduce_dims=reduce_dims,
                preserve_dims=preserve_dims,
                weights=weights,
            )
        self.left = pit_cdf["left"]
        self.right = pit_cdf["right"]
        self.pit_uniform_endpoints = pit_cdf["pit_uniform_endpoints"]

    def plotting_points(self) -> XarrayLike:
        """
        Returns an xarray object with the plotting points for the graph of the
        PIT CDF :math:`F`. The x (horizontal) plotting positions are values from the "pit_x_value" index, while the
        y (vertical) plotting positions are values in the output xarray object.
        All other points on the graph are determined by linear interpolation.

        Also useful for generating PIT-uniform probability plots.

        Note that coordinates in the "pit_x_value" dimension will have duplicate values.
        For a parametric approach to plotting points without duplicate coordinate values,
        see the ``plotting_points_parametric`` method.
        """
        return _get_plotting_points(self.left, self.right)

    def plotting_points_parametric(self) -> dict:
        """
        Returns the plotting points for the graph of CDF :math:`F` of the PIT distribution.
        Points on the plot are given by :math:`(x(t),y(t))` for plotting parameter :math:`t`,
        with any gaps determined by linear interpolation.

        Also useful for generating PIT-uniform probability plots.

        Returns:
            dictionary of xarray objects with the following keys:
            - "x_plotting_position", containing :math:`x(t)` values
            - "y_plotting_position", containing :math:`y(t)` values
            with :math:`t` values in the "plotting_point" dimension.
        """
        return _get_plotting_points_param(self.left, self.right)

    def hist_values(self, bins: int, right: bool = True) -> XarrayLike:
        """
        Returns an xarray object with the PIT histogram values across appropriate dimensions,
        with additional coordinates "bin_left_endpoint", "bin_right_endpoint" and "bin_centre".

        Args:
            bins: the number of bins in the histogram.
            right: If True, histogram bins always include the rightmost edge. If False,
                bins always include the leftmost edge.

        Example:
            If ``bins=5`` and ``right=True`` then the bins are the set of intervals
            :math:`\\{[0, 0.2], (0.2, 0.4], (0.4, 0.6], (0.6, 0.8], (0.8, 1]\\}`.

        See also:
            - :py:func:`scores.probability.rank_histogram`
        """
        # calculating _pit_hist_right or _pit_hist_left does not work with chunks,
        # because dask wants to chunk over the 'pit_x_value' dimension, but
        # _pit_hist_right and _pit_hist_left merely wants to sample a small number of points
        # from the 'pit_x_value'. So compute chunks at this stage.
        self.left = self.left.compute()
        self.right = self.right.compute()

        if right:
            return _pit_hist_right(self.left, self.right, bins)
        return _pit_hist_left(self.left, self.right, bins)

    def alpha_score(self, negative_orientation: bool = True) -> XarrayLike:
        """
        Returns the 'alpha score' (Renard, et. al., 2010), which is a measure of how close the
        PIT distribution :math:`F` is to the uniform distribution on the closed unit interval :math:`[0,1]`.
        When ``negative_orientation=True``, the formula is for the alpha score is
            :math:`\\int_0^1 |F(x) - x|\\,\\text{d}x,`
        so that the lower the score, the closer :math:`F` is to the uniform distribution.
        In this case the alpha score takes values between 0 and 0.5.
        When ``negative_orientation=False``, the formula is for the alpha score is
            :math:`1 - 2 \\int_0^1 |F(x) - x|\\,\\text{d}x,`
        so that the higher the score, the closer :math:`F` is to the uniform distribution.
        In this case the alpha score takes values between 0 and 1.

        References:
            - Renard, B., Kavetski, D., Kuczera, G., Thyer, M., & Franks, S. W. (2010). \
                Understanding predictive uncertainty in hydrologic modeling: \
                The challenge of identifying input and structural errors. \
                Water Resources Research, 46(5). https://doi.org/10.1029/2009WR008328
        """
        return _alpha_score(self.left, self.right, negative_orientation=negative_orientation)

    def expected_value(self) -> XarrayLike:
        """
        Returns the expected value of the PIT distribution :math:`F`.
        An expected value greater than 0.5 indicates an under-prediction tendency.
        An expected value less than 0.5 indicates an over-prediction tendency.
        """
        return _expected_value(self.plotting_points())

    def variance(self) -> XarrayLike:
        """
        Returns the variance of the PIT distribution :math:`F`.
        A variance greater than 1/12 indicates predictive under-dispersion.
        A variance less than 1/12 indicates predictive over-dispersion.
        """
        return _variance(self.plotting_points())


class PitFcstAtObs:
    """
    Calculates the probability integral transform (PIT) for a set of forecast cases and
    corresponding observations. The calculated PIT can be a (possibly weighted) average
    over specified dimensions. Calculations are performed and interpreted as follows.

    Given a predictive cumulative distribution function (CDF) :math:`G` and corresponding
    observation :math:`y`, the corresponding PIT value is a uniform distribution on the
    closed interval :math:`[G(y-), G(y)]`, where :math:`G(y-)` denotes the left-hand
    limit of :math:`G` at :math:`y`. This is the most general form of PIT and handles cases
    where the observation :math:`y` coincides with a point of discontinuity of a predictive CDF
    :math:`G`. See Gneiting and Ranjan (2013) and Taggart (2022).

    Weighted means of PIT values are simply weighted means of CDFs of those uniform distributions
    on the closed interval :math:`[G(y-), G(y)]`. The weighted mean is itself a CDF, here denoted :math:`F`
    and called the PIT distribution for the set of forecast cases.
    Left-hand and right-hand limits of :math:`F` are accessible via the attributes ``.left`` and ``.right``.
    Important statistics related to the PIT distribution :math:`F`, such as bar heights of PIT histograms,
    the graph of :math:`F`, plotting points for PIT uniform probability plots, the alpha score, and
    expected value and variance of :math:`F`, can be calculated from :math:`F` and are accessible via
    methods on the ``PitFcstAtObs`` class.

    In the ``PitFcstAtObs`` implementation of PIT, the user supplies the values :math:`G(y)` and
    (optionally) :math:`G(y-)`. If calculating ``PitFcstAtObs()`` is slow or runs into memory
    issues, then rounding the values of :math:`G(y)` and :math:`G(y-)` may resolve this problem with
    very little change to calculated statistics. If the predictive distributions are in the form of an ensemble
    or of CDFs whose values are harder to evaluate at the observations, consider using the
    ``Pit`` implementation instead.

    Args:
        fcst_at_obs: an xarray object of values of the forecast CDF evaluated at each
            corresponding observation.
        fcst_at_obs_left: an xarray object of left-hand limits of the forecast CDF at each
            corresponding observation. Only needs to be supplied if there are any cases
            where the left-hand limit does not equal the value at the forecast CDF.
        reduce_dims: Optionally specify which dimensions to reduce when calculating the
            PIT CDF values, where the mean is taken over all forecast cases.
            All other dimensions will be preserved. As a special case, 'all' will allow
            all dimensions to be reduced. Only one of ``reduce_dims`` and ``preserve_dims``
            can be supplied. The default behaviour if neither are supplied is to reduce all dims.
        preserve_dims: Optionally specify which dimensions to preserve when calculating the
            PIT CDF values, where the mean is taken over all forecast cases.
            All other dimensions will be reduced. As a special case, 'all' will allow
            all dimensions to be preserved. Only one of ``reduce_dims`` and ``preserve_dims``
            can be supplied. The default behaviour if neither are supplied is to reduce all dims.
        weights: Optionally provide an array for weighted averaging (e.g. by area, by latitude,
            by population, custom) of PIT CDF values across all forecast cases.

    Attributes:
        left: xarray object giving the left-hand limits :math:`F(x-)` of the CDF :math:`F` of the
            (possibly weighted) PIT distribution across all forecast cases. The points :math:`x`
            are in the dimension "pit_x_value". The CDF :math:`F` of the PIT distribution is completely
            reconstructable using values from ``.left``, ``.right`` and linear interpolation.
        right: xarray object giving the values :math:`F(x)` of the CDF :math:`F` of the
            (possibly weighted) PIT distribution across all forecast cases. The points :math:`x`
            are in the dimension "pit_x_value". The CDF :math:`F` of the PIT distribution is completely
            reconstructable using values from ``.left``, ``.right`` and linear interpolation.
        pit_uniform_endpoints: xarray object giving the interval :math:`[G(y-), G(y)]`
            of the uniform distribution generated from the probability integral transform for each
            forecast--observation pair :math:`(G,y)`. Dimensions are not reduced and weights are ignored.
            The left and right endpoints of the interval are given by the "lower" and "upper" coordinates
            in the "uniform_endpoint" dimension.

    Raises:
        ValueError: when dimensions of ``fcst``, ``obs`` or ``weights`` contain any of the following reserved names:
            "uniform_endpoint", "pit_x_value", "x_plotting_position", "y_plotting_position", "plotting_point",
            "bin_left_endpoint", "bin_right_endpoint", "bin_centre",
        ValueError: when ``fcst_at_obs_left`` is not ``None``, and ``fcst_at_obs`` and ``fcst_at_obs_left``
            do not have identical shape, dimensions and coordinates.
        ValueError: when any values of ``fcst`` are less then 0 or greater than 1.
        ValueError: when ``fcst_at_obs_left`` is not ``None`` and
            any values of ``fcst_at_obs_left`` are less then 0 or greater than 1.
        ValueError: when ``fcst_at_obs_left`` is not ``None`` and
            any values of ``fcst_at_obs_left`` are greater than ``fcst``.

    References:
        - Gneiting, T., & Ranjan, R. (2013). Combining predictive distributions. Electron. J. Statist. 7: 1747-1782 \
            https://doi.org/10.1214/13-EJS823
        - Taggart, R. J. (2022). Assessing calibration when predictive distributions have discontinuities. \
            Bureau Research Report 64, https://nla.gov.au/nla.obj-3079961862/view

    See also:
            - :py:func:`scores.probability.Pit`
            - :py:func:`scores.probability.rank_histogram`

    Examples:
        Calculate the PIT for an under-dispersive forecast, and calculate various diagnostics and statistics.

        >>> import xarray as xr
        >>> from scipy.stats import norm
        >>> from scores.probability import PitFcstAtObs
        >>> # observations generated from a normal distribution with
        >>> # mean 0 and standard deviation 2
        >>> obs = xr.DataArray(norm.rvs(scale=2, size=(500)), dims=['time'])
        >>> # forecasts are normal distributions with mean 0 and standard deviation 1
        >>> # evaluate the forecast CDFs at the observations
        >>> fcst_at_obs = xr.DataArray(norm.cdf(obs), dims=['time'])
        >>> pit = PitFcstAtObs(fcst_at_obs)
        >>> # bar heights for a PIT histogram
        >>> histogram_values = pit.hist_values(10)
        >>> # plot the CDF of the PIT distribution
        >>> pit.plotting_points().plot()
        >>> # the expected value of the PIT distribution
        >>> pit_ev = pit.expected_value()
        >>> # the variance of the PIT distribution
        >>> pit_var = pit.variance()
        >>> # the alpha score of the PIT distribution
        >>> pit_alpha = pit.alpha_score()

        Other examples are found in the ``scores`` tutorial for PIT.
    """

    def __init__(
        self,
        fcst_at_obs: XarrayLike,
        *,  # Force keywords arguments to be keyword-only
        fcst_at_obs_left: Optional[XarrayLike] = None,
        reduce_dims: Optional[FlexibleDimensionTypes] = None,
        preserve_dims: Optional[FlexibleDimensionTypes] = None,
        weights: Optional[XarrayLike] = None,
    ):
        """
        Calculates the left-hand and right-hand limits of the PIT distribution :math:`F`,
        accessible via the attributes ``.left`` and ``.right``, as well as the interval
        :math:`[G(y-),G(y)]` of each probability integral transformed forecast--observation pair
        :math:`(G,y)`, accessible via the attribute ``.pit_uniform_endpoints``.
        """
        pit_cdf = _pit_values_for_fcst_at_obs(fcst_at_obs, fcst_at_obs_left, reduce_dims, preserve_dims, weights)
        self.left = pit_cdf["left"]
        self.right = pit_cdf["right"]
        self.pit_uniform_endpoints = pit_cdf["pit_uniform_endpoints"]

    def plotting_points(self) -> XarrayLike:
        """
        Returns an xarray object with the plotting points for the graph of the
        PIT CDF :math:`F`. The x (horizontal) plotting positions are values from the "pit_x_value" index,
        while the y (vertical) plotting positions are values in the output xarray object.
        All other points on the graph are determined by linear interpolation.

        Also useful for generating PIT-uniform probability plots.

        Note that coordinates in the "pit_x_value" dimension will have duplicate values.
        For a parametric approach to plotting points without duplicate coordinate values,
        see the ``plotting_points_parametric`` method.
        """
        return _get_plotting_points(self.left, self.right)

    def plotting_points_parametric(self) -> dict:
        """
        Returns the plotting points for the graph of CDF :math:`F` of the PIT distribution.
        Points on the plot are given by :math:`(x(t),y(t))` for plotting parameter :math:`t`,
        with any gaps determined by linear interpolation.

        Also useful for generating PIT-uniform probability plots.

        Returns:
            dictionary of xarray objects with the following keys:
            - "x_plotting_position", containing :math:`x(t)` values
            - "y_plotting_position", containing :math:`y(t)` values
            with :math:`t` values in the "plotting_point" dimension.
        """
        return _get_plotting_points_param(self.left, self.right)

    def hist_values(self, bins: int, right: bool = True) -> XarrayLike:
        """
        Returns an xarray object with the PIT histogram values across appropriate dimensions,
        with additional coordinates "bin_left_endpoint", "bin_right_endpoint" and "bin_centre".

        Args:
            bins: the number of bins in the histogram.
            right: If True, histogram bins always include the rightmost edge. If False,
                bins always include the leftmost edge.

        Example:
            If ``bins=5`` and ``right=True`` then the bins are the set of intervals
            :math:`\\{[0, 0.2], (0.2, 0.4], (0.4, 0.6], (0.6, 0.8], (0.8,1]\\}`.

        See also:
            - :py:func:`scores.probability.rank_histogram`
        """
        # calculating _pit_hist_right or _pit_hist_left does not work with chunks,
        # because dask wants to chunk over the 'pit_x_value' dimension, but
        # _pit_hist_right and _pit_hist_left merely wants to sample a small number of points
        # from the 'pit_x_value'. So compute chunks at this stage.
        if self.left.chunks is not None:
            self.left = self.left.compute()
            self.right = self.right.compute()

        if right:
            return _pit_hist_right(self.left, self.right, bins)
        return _pit_hist_left(self.left, self.right, bins)

    def alpha_score(self, negative_orientation: bool = True) -> XarrayLike:
        """
        Returns the 'alpha score' (Renard, et. al., 2010), which is a measure of how close the
        PIT distribution :math:`F` is to the uniform distribution on the closed unit interval :math:`[0,1]`.
        When ``negative_orientation=True``, the formula is for the alpha score is
            :math:`\\int_0^1 |F(x) - x|\\,\\text{d}x,`
        so that the lower the score, the closer :math:`F` is to the uniform distribution.
        In this case the alpha score takes values between 0 and 0.5.
        When ``negative_orientation=False``, the formula is for the alpha score is
            :math:`1 - 2 \\int_0^1 |F(x) - x|\\,\\text{d}x,`
        so that the higher the score, the closer :math:`F` is to the uniform distribution.
        In this case the alpha score takes values between 0 and 1.

        References:
            - Renard, B., Kavetski, D., Kuczera, G., Thyer, M., & Franks, S. W. (2010). \
                Understanding predictive uncertainty in hydrologic modeling: \
                The challenge of identifying input and structural errors. \
                Water Resources Research, 46(5). https://doi.org/10.1029/2009WR008328
        """
        return _alpha_score(self.left, self.right, negative_orientation=negative_orientation)

    def expected_value(self) -> XarrayLike:
        """
        Returns the expected value of the PIT distribution :math:`F`.
        An expected value greater than 0.5 indicates an under-prediction tendency.
        An expected value less than 0.5 indicates an over-prediction tendency.
        """
        return _expected_value(self.plotting_points())

    def variance(self) -> XarrayLike:
        """
        Returns the variance of the PIT distribution :math:`F`.
        A variance greater than 1/12 indicates predictive under-dispersion.
        A variance less than 1/12 indicates predictive over-dispersion.
        """
        return _variance(self.plotting_points())


#####################################################
# Functions for calculating __init__ of both classes
#####################################################


def _dims_for_mean_with_checks(
    fcst: XarrayLike,
    obs: XarrayLike,
    special_fcst_dim: Optional[str],
    weights: Optional[XarrayLike],
    reduce_dims: Optional[FlexibleDimensionTypes],
    preserve_dims: Optional[FlexibleDimensionTypes],
) -> set[Hashable]:
    """
    Given inputs for Pit or PitFcstAtObs, checks that XarrayLike inputs don't use
    RESERVES_NAMES and then gathers dimensions, returning the set of dimensions for
    calculating the mean.

    When applying to `PitFcstAtObs` inputs, use `fcst=obs=fcst_at_obs`.
    """
    all_dims = set(fcst.dims).union(obs.dims)

    weights_dims = None
    if weights is not None:
        weights_dims = weights.dims
        all_dims = all_dims.union(weights_dims)

    if len([dim for dim in RESERVED_NAMES if dim in all_dims]) > 0:
        raise ValueError(
            f"The following names are reserved and should not be among the dimensions of \
            xarray inputs: {', '.join(RESERVED_NAMES)}"
        )

    dims_for_mean = gather_dimensions(
        fcst.dims,
        obs.dims,
        weights_dims=weights_dims,
        reduce_dims=reduce_dims,
        preserve_dims=preserve_dims,
        score_specific_fcst_dims=special_fcst_dim,
    )
    return dims_for_mean


def _right_left_checks(
    right: XarrayLike, left: Optional[XarrayLike], threshold_dim: Optional[str], right_arg_name: str, left_arg_name: str
):
    """
    Raises:
        - ValueError if `right` or `left` have values less than 0 or greater than 1
        - ValueError if `right[threshold_dim]` is not strictly increasing
        - ValueError if, when `left` is not None, `right` and `left` don't have the same shape, dims or coords
        - ValueError if, when `left` is not None, `right < left` at some point
    """
    # first check right
    if isinstance(right, xr.DataArray):
        if not cdf_values_within_bounds(right):
            raise ValueError(f"`{right_arg_name}` values must be between 0 and 1 inclusive.")
    else:
        for var in right.data_vars:
            if not cdf_values_within_bounds(right[var]):
                raise ValueError(f"`{right_arg_name}` values must be between 0 and 1 inclusive.")
    if (threshold_dim is not None) and (not coords_increasing(right, threshold_dim)):
        raise ValueError(f"coordinates along `{right_arg_name}[threshold_dim]` are not strictly increasing")

    # then check left if appropriate
    if left is not None:
        # check that both fcst and fcst_left have same shape, coords and dims
        if not xr.ones_like(right).equals(xr.ones_like(left)):
            raise ValueError(
                f"If `{left_arg_name}` is not `None`, `{right_arg_name}` and `{left_arg_name}` "
                "must have same shape, dimensions and coordinates"
            )
        if isinstance(right, xr.DataArray):
            if not cdf_values_within_bounds(left):
                raise ValueError(f"`{left_arg_name}` values must be between 0 and 1 inclusive.")
            if (left > right).any():
                raise ValueError(f"`{left_arg_name}` must not exceed `{right_arg_name}`")
        else:
            for var in left.data_vars:
                if not cdf_values_within_bounds(left[var]):
                    raise ValueError(f"`{left_arg_name}` values must be between 0 and 1 inclusive.")
                if (left[var] > right[var]).any():
                    raise ValueError(f"`{left_arg_name}` must not exceed `{right_arg_name}`")


def _pit_values_for_fcst_at_obs(
    fcst_at_obs: XarrayLike,
    fcst_at_obs_left: Optional[XarrayLike],
    reduce_dims: Optional[FlexibleDimensionTypes],
    preserve_dims: Optional[FlexibleDimensionTypes],
    weights: Optional[XarrayLike],
) -> dict:
    """
    A private function to compute `PitFcstAtObs.__init__`.
    Returns `PitFcstAtObs().left`, `PitFcstAtObs().right` and `PitFcstAtObs().pit_uniform_endpoints
    in the form of a dictionary with keys 'left', 'right' and 'pit_uniform_endpoints'.

    See docstring `PitFcstAtObs.__init__` for details.
    """
    _right_left_checks(fcst_at_obs, fcst_at_obs_left, None, "fcst_at_obs", "fcst_at_obs_left")

    dims_for_mean = _dims_for_mean_with_checks(fcst_at_obs, fcst_at_obs, None, weights, reduce_dims, preserve_dims)

    if fcst_at_obs_left is None:
        fcst_at_obs_left = fcst_at_obs.copy()
    else:
        fcst_at_obs, fcst_at_obs_left = broadcast_and_match_nan(fcst_at_obs, fcst_at_obs_left)

    pit_values = xr.concat(
        [
            fcst_at_obs_left.assign_coords(uniform_endpoint="lower").expand_dims("uniform_endpoint"),
            fcst_at_obs.assign_coords(uniform_endpoint="upper").expand_dims("uniform_endpoint"),
        ],
        "uniform_endpoint",
    )

    # convert to CDF format, take weighted means, and output dictionary
    result = _pit_values_final_processing(pit_values, weights, dims_for_mean)
    result["pit_uniform_endpoints"] = pit_values

    return result


def _pit_values_final_processing(
    pit_values: XarrayLike, weights: Optional[XarrayLike], dims_for_mean: set[Hashable]
) -> dict:
    """
    Given PIT values in the format [lower,upper], representing the uniform distribution
    on the closed interval [lower,upper], converts to CDF format, with output 'left' giving
    the left limit of the CDF and 'right' giving the value at the CDF. Weighted means are then
    taken across all the CDFs. The output is scaled so that it is in CDF format (i.e., max value is 1)

    Args:
        pit_values: values of the PIT in [upper, lower] format, including dimension 'uniform_endpoint'
            that has two coordinates "upper" and "lower"
        weights: optional array of weights when calculating the weighted mean
        dims_for_mean: list, set etc of dimensions over which to apply the mean.

    Returns:
        dictionary of three xarray objects, with keys 'left', 'right' and 'pit_uniform_endpoints'.
        These contain values of the left-hand and right-hand limits of the mean weighted PIT in CDF format,
        and the `pit_values` input to this function.
    """
    # convert to F(x-), F(x) format
    pit_cdf = _pit_cdfvalues(pit_values)

    pit_cdf_left = aggregate(pit_cdf["left"], reduce_dims=dims_for_mean, weights=weights)
    pit_cdf_right = aggregate(pit_cdf["right"], reduce_dims=dims_for_mean, weights=weights)

    # rescale CDFs so that their max value is 1.
    # This corrects for weights that don't sum to 1.
    cdf_right_max = pit_cdf_right.max("pit_x_value")
    valid_max = cdf_right_max.notnull() & (cdf_right_max > 0)
    safe_cdf_right_max = xr.where(valid_max, cdf_right_max, 1)
    pit_cdf_right = pit_cdf_right / safe_cdf_right_max
    pit_cdf_left = pit_cdf_left / safe_cdf_right_max
    pit_cdf_right = pit_cdf_right.where(valid_max)
    pit_cdf_left = pit_cdf_left.where(valid_max)

    return {"left": pit_cdf_left, "right": pit_cdf_right, "pit_uniform_endpoints": pit_values}


def _pit_values_for_ens(fcst: XarrayLike, obs: XarrayLike, ens_member_dim: str) -> XarrayLike:
    """
    For each forecast case in the form of an ensemble, the PIT value of the ensemble for
    the corresponding observation is a uniform distribution over the closed interval
    [lower,upper], where
        lower = (count of ensemble members strictly less than the observation) / n
        upper = (count of ensemble members not exceeding the observation) / n
        n = size of the ensemble.

    Returns an array of [lower,upper] values in the dimension 'uniform_endpoint'.

    Args:
        fcst: array of forecast values, including dimension `ens_member_dim`
        obs: array of forecast values, excluding `ens_member_dim`
        ens_member_dim: name of the ensemble member dimension in `fcst`

    Returns:
        array of PIT values in the form [lower,upper], with dimensions
        'uniform_endpoint', all dimensions in `obs` and all dimensions in `fcst`
        excluding `ens_member_dim`.
    """
    ensemble_size = fcst.count(ens_member_dim)
    obs_valid = obs.notnull()
    divisor_valid = obs_valid & (ensemble_size > 0)
    safe_ensemble_size = xr.where(divisor_valid, ensemble_size, 1)
    pit_lower = ((fcst < obs).sum(ens_member_dim) / safe_ensemble_size).where(divisor_valid)
    pit_upper = ((fcst <= obs).sum(ens_member_dim) / safe_ensemble_size).where(divisor_valid)

    pit_lower = pit_lower.assign_coords(uniform_endpoint="lower").expand_dims("uniform_endpoint")
    pit_upper = pit_upper.assign_coords(uniform_endpoint="upper").expand_dims("uniform_endpoint")

    return xr.concat([pit_lower, pit_upper], "uniform_endpoint")


def _pit_distribution_for_cdf(
    fcst: XarrayLike,
    obs: XarrayLike,
    threshold_dim: str,
    *,  # Force keywords arguments to be keyword-only
    fcst_left: Optional[XarrayLike] = None,
    reduce_dims: Optional[FlexibleDimensionTypes] = None,
    preserve_dims: Optional[FlexibleDimensionTypes] = None,
    weights: Optional[XarrayLike] = None,
) -> dict:
    """
    Each forecast case is an array of cumulative distribution function (CDF) values :math:`G(x)`,
    which is the probability that the random variable being forecast does not exceed :math:`x`.
    Given any observation :math:`y`, the probability integral transform (PIT) value at :math:`G`
    is a uniform distribution on the closed interval :math:`[G(y-), G(y)]`, where :math:`G(y-)`
    denotes the left-hand limit of :math:`G` at :math:`y`.

    This function outputs values of PIT by representing each uniform distribution on
    :math:`[G(y-), G(y)]` as the corresponding uniform CDF :math:`F`. We call :math:`F`
    the 'PIT distribution' for the forecast--observation pair :math:`(G,y)`. Values :math:`F(x-)`
    and :math:`F(x)` are given for an optimal set of points :math:`x` satisfying
    :math:`0 <= x <= 1`. The set is optimal in the sense that the value of :math:`F`
    elsewhere can be determined via linear interpolation, whilst no smaller set of values
    has this property.

    Dimensions are be reduced by taking (possibly weighted) means of the PIT distribution
    :math:`F(x)`. The weighted means can also be interpreted a values of CDFs, and exact
    intermediate values also attained via linear interpolation.

    Args:
        fcst: xarray object of CDF forecast values, containing the dimension `threshold_dim`,
            or a dictionary containing the keys "left" and "right", with corresponding values
            xarray objects giving the left-hand and right-hand limits of the fcst CDF
            along the dimension `threshold_dim`.
        obs: an xarray object of observations.
        threshold_dim: name of the threshold dimension in ``fcst``, such that the probability
            of not exceeding a particular threshold is one of the corresponding values of ``fcst``.
        fcst_left: xarray object of forecast CDF left-handed limit values. If None, it is
            assumed that the forecast CDF is continuous.
        reduce_dims: Optionally specify which dimensions to reduce when calculating the
            PIT distribution, where the mean is taken over all forecast cases.
            All other dimensions will be preserved. As a special case, 'all' will allow
            all dimensions to be reduced. Only one of ``reduce_dims`` and ``preserve_dims``
            can be supplied. The default behaviour if neither are supplied is to reduce all dims.
        preserve_dims: Optionally specify which dimensions to preserve when calculating the
            PIT distribution, where the mean is taken over all forecast cases.
            All other dimensions will be reduced. As a special case, 'all' will allow
            all dimensions to be preserved, apart from ``severity_dim`` and ``prob_threshold_dim``.
            Only one of ``reduce_dims`` and ``preserve_dims`` can be supplied. The default
            behaviour if neither are supplied is to reduce all dims.
        weights: Optionally provide an array for weighted averaging (e.g. by area, by latitude,
            by population, custom) of PIT CDF values across all forecast cases.

    Returns:
        a dictionary with the following keys and values:
        - "left": an xarray object containing the left-hand limits :math:`F(x-)` of the PIT CDF values
        - "right": an xarray object containing the values :math:`F(x)` of the PIT CDF values
        - "pit_uniform_endpoints": an xarray object containing the values :math:`[G(y-), G(y)]`,
            with lower and upper endpoints indexed by the dimension "uniform_endpoint"

    Raises:
        - ValueError if dimensions of ``fcst``, ``obs`` or ``weights`` contain any of the following reserved names:
            'uniform_endpoint', 'pit_x_value', 'x_plotting_position', 'y_plotting_position', 'plotting_point'
        - ValueError if, when ``fcst`` is a dictionary, ``fcst['left']`` and ``fcst['right']`` do not have the same
            shape, dimensions and coordinates.
        - ValueError if values in ``fcst`` arrays or sets are less than 0 or greater than 1, ignoring NaNs.
        - ValueError if coordinates in the ``fcst`` ``threshold_dim`` dimension are not increasing.

    Warns:
        - if any values in ``obs`` lie outside the range of values in the forecast ``threshold_dim`` dimension.
        - if any forecast values have NaN
    """
    _right_left_checks(fcst, fcst_left, threshold_dim, "fcst", "fcst_left")

    dims_for_mean = _dims_for_mean_with_checks(fcst, obs, threshold_dim, weights, reduce_dims, preserve_dims)

    if fcst_left is None:
        fcst_left = fcst.copy()

    # PIT values in [G(y-), G(y)] format
    pit_values = _pit_values_for_cdf(fcst_left, fcst, obs, threshold_dim)

    # convert to CDF format, take weighted means, and output dictionary
    result = _pit_values_final_processing(pit_values, weights, dims_for_mean)

    return result


def _pit_values_for_cdf_array(
    fcst_left: xr.DataArray, fcst_right: xr.DataArray, obs: xr.DataArray, threshold_dim: str
) -> xr.DataArray:
    """
    For each forecast case in the form of an CDF G (in xr.DataArray format), the PIT value of G
    for the corresponding observation y is a uniform distribution over the closed
    interval [lower,upper], where
        lower = G(y-)
        upper = G(y)
    and G(y-) denotes the left-hand limit of F at y.

    Returns an array of [lower,upper] values in the dimension 'uniform_endpoint'.

    It is assumed that `fcst_left` and `fcst_right` have the same shape, dims and coords.

    Args:
        fcst_left: array of forecast CDF left-limit values, including dimension `threshold_dim`
        fcst_right: array of forecast CDF right-limit values, including dimension `threshold_dim`.
            Assumed to have same shape and dimensions at `fcst_left`
        obs: array of observed values; does not have dimension `threshold_dim`
        threshold_dim: name of the threshold dimension in `fcst`

    Returns:
        array of PIT values in the form [lower,upper], with dimensions
        'uniform_endpoint', all dimensions in `obs` and all dimensions in `fcst`
        excluding `ens_member_dim`.
    """
    # check whether observations are in the threshold_dim range, if not issue a warning
    max_obs = obs.max()
    min_obs = obs.min()
    max_thld = fcst_left[threshold_dim].max()
    min_thld = fcst_left[threshold_dim].min()

    if max_obs > max_thld:
        warnings.warn(
            "Some observations were greater than the maximum `threshold_dim` value in your `fcst`. "
            "The value of the fcst CDF at these observations will be set to 1.",
            UserWarning,
        )
    if min_obs < min_thld:
        warnings.warn(
            "Some observations were less than the minimum `threshold_dim` value in your `fcst`. "
            "The value of the fcst CDF at these observations will be set to 0.",
            UserWarning,
        )

    # check whether fcst values are NaN, if so issue a warning
    if bool(fcst_left.isnull().any()) or bool(fcst_right.isnull().any()):
        warnings.warn(
            "Some forecast CDF values are NaN. In such cases, the entire forecast CDF will be treated as NaN. "
            "To avoid this, you can fill NaNs using `scores.processing.cdf.fill_cdf`.",
            UserWarning,
        )

    flatten_obs = np.unique(obs.values.flatten())
    flatten_obs = flatten_obs[~np.isnan(flatten_obs)]

    # classify the observations
    obs_in_thlds = [x for x in flatten_obs if x in fcst_right[threshold_dim]]
    obs_gt_thlds = [x for x in flatten_obs if x > max_thld]
    obs_lt_thlds = np.array([x for x in flatten_obs if x < min_thld])
    obs_between_thlds = [
        x for x in flatten_obs if x not in set(obs_in_thlds).union(set(obs_gt_thlds)).union(set(obs_lt_thlds))
    ]

    # extend the fcst CDF values so they can be evaluated at obs
    fcst_left_at_obs = []
    fcst_right_at_obs = []

    # start with the original forecast
    if len(obs_in_thlds) > 0:
        fcst_left_at_obs.append(fcst_left.sel({threshold_dim: obs_in_thlds}))
        fcst_right_at_obs.append(fcst_right.sel({threshold_dim: obs_in_thlds}))

    # calculate the interpolated values using interp1d
    # (xarray interpolate_na not suitable as it cannot handle duplicate x values)
    if len(obs_between_thlds) > 0:
        plotting_points = xr.concat([fcst_left, fcst_right], threshold_dim).sortby(threshold_dim)
        threshold_axis = plotting_points.get_axis_num(threshold_dim)
        y_values = plotting_points.values
        x_values = plotting_points[threshold_dim].values

        interpolated_values = interp1d(x_values, y_values, axis=threshold_axis, kind="linear")(obs_between_thlds)

        # turn the numpy array into an xarray array
        interpolated_coords = dict(fcst_left.coords)
        interpolated_coords[threshold_dim] = obs_between_thlds
        interpolated_values = xr.DataArray(data=interpolated_values, dims=fcst_left.dims, coords=interpolated_coords)

        fcst_left_at_obs.append(interpolated_values)
        fcst_right_at_obs.append(interpolated_values)

    # add the extrapolated values, with values of either 0 or 1
    if len(obs_gt_thlds) > 0:
        high_thld = xr.DataArray(data=obs_gt_thlds, dims=[threshold_dim], coords={threshold_dim: obs_gt_thlds})
        high_thld = xr.broadcast(high_thld, fcst_left)[0].sel({threshold_dim: obs_gt_thlds})
        high_thld = xr.ones_like(high_thld)
        fcst_left_at_obs.append(high_thld)
        fcst_right_at_obs.append(high_thld)

    if len(obs_lt_thlds) > 0:
        low_thld = xr.DataArray(data=obs_lt_thlds, dims=[threshold_dim], coords={threshold_dim: obs_lt_thlds})
        low_thld = xr.broadcast(low_thld, fcst_left)[0].sel({threshold_dim: obs_lt_thlds})
        low_thld = xr.zeros_like(low_thld)
        fcst_left_at_obs.append(low_thld)
        fcst_right_at_obs.append(low_thld)

    # combine the data
    no_nans = (fcst_right.notnull() & fcst_left.notnull()).all(threshold_dim)
    fcst_left_at_obs = xr.concat(fcst_left_at_obs, threshold_dim)
    fcst_right_at_obs = xr.concat(fcst_right_at_obs, threshold_dim)

    pit_cdf_left = (
        fcst_left_at_obs.where(fcst_left_at_obs[threshold_dim] == obs)
        .mean(threshold_dim)
        .assign_coords(uniform_endpoint="lower")
        .expand_dims("uniform_endpoint")
        .where(no_nans)
    )
    pit_cdf_right = (
        fcst_right_at_obs.where(fcst_right_at_obs[threshold_dim] == obs)
        .mean(threshold_dim)
        .assign_coords(uniform_endpoint="upper")
        .expand_dims("uniform_endpoint")
        .where(no_nans)
    )
    result = xr.concat([pit_cdf_left, pit_cdf_right], "uniform_endpoint")

    return result


def _pit_values_for_cdf(
    fcst_left: XarrayLike, fcst_right: XarrayLike, obs: XarrayLike, threshold_dim: str
) -> xr.Dataset:
    """
    Does the same as `_pit_values_for_cdf_array` but supports XarrayLike inputs.

    Args:
        fcst_left: xarray object forecast CDF left-limit values, including dimension `threshold_dim`
        fcst_right: xarray object forecast CDF right-limit values, including dimension `threshold_dim`.
            Assumed to have same shape, variables, coords, etc as fcst_right
        obs: array of forecast values, excluding `threshold_dim`
        threshold_dim: name of the threshold dimension in `fcst`

    Returns:
        array of PIT values in the form [lower,upper], with dimensions
        'uniform_endpoint', all dimensions in `obs` and all dimensions in `fcst`
        excluding `ens_member_dim`.
    """
    if isinstance(fcst_left, xr.DataArray) and isinstance(obs, xr.DataArray):
        return _pit_values_for_cdf_array(fcst_left, fcst_right, obs, threshold_dim)
    elif isinstance(fcst_left, xr.Dataset) and isinstance(obs, xr.DataArray):
        return xr.merge(
            [
                _pit_values_for_cdf_array(fcst_left[var], fcst_right[var], obs, threshold_dim).rename(var)
                for var in fcst_left.data_vars
            ]
        )
    elif isinstance(fcst_left, xr.DataArray):
        return xr.merge(
            [
                _pit_values_for_cdf_array(fcst_left, fcst_right, obs[var], threshold_dim).rename(var)
                for var in obs.data_vars
            ]
        )
    else:
        return xr.merge(
            [
                _pit_values_for_cdf_array(fcst_left[var], fcst_right[var], obs[var], threshold_dim).rename(var)
                for var in obs.data_vars
            ]
        )


def _get_pit_x_values(pit_values: XarrayLike) -> xr.DataArray:
    """
    Returns a data array of consisting of exactly those x-axis values needed for
    constructing a uniform PIT probability plot for the given array of `pit_values`.

    Args:
        pit_values: output from `_pit_values_for_ens` or `_pit_values_for_cdf`, containing
            dimension 'uniform_endpoint' with coordinates ['lower', 'upper']

    Returns:
        xarray data array of x-axis values, indexed by the same values on the dimension
        'pit_x_value'
    """
    if isinstance(pit_values, xr.Dataset):
        pit_values = pit_values.to_dataarray()

    x_values = np.unique(pit_values)
    x_values = np.unique(np.concatenate((np.array([0.0, 1.0]), x_values)))
    x_values = x_values[~np.isnan(x_values)]
    x_values = xr.DataArray(data=x_values, dims=["pit_x_value"], coords={"pit_x_value": x_values})
    return x_values


def _pit_distribution_for_jumps(pit_values: XarrayLike, x_values: xr.DataArray) -> dict:
    """
    Gives the values F(x) where F is the CDF for a pit value,
    and x comes from `x_values`, given that the CDF jumps at x. This occurs precisely
    when upper == lower in the [lower, upper] representation of the PIT value.
    If this condition fails, NaNs are returned.

    It is assumed that `x_values` contains all the values in `pit_values`.

    Args:
        pit_values: output from `_pit_values_for_ens` or `_pit_values_for_cdf`, containing
            dimension 'uniform_endpoint' with coordinates ['lower', 'upper']
        x_values: xr.DataArray of x values, with dimension 'pit_x_value', output from
            `_get_pit_x_values`

    Returns:
        dictionary of cdf values, with keys 'left' and 'right',
        representing the left and right hand limits of the cdf values F(x).
        Each value in the dictionary is an xarray object representing limits of F(x),
        where x is represented by values in the 'pit_x_value' dimension.
    """
    lower_values = pit_values.sel(uniform_endpoint="lower").drop_vars("uniform_endpoint")
    upper_values = pit_values.sel(uniform_endpoint="upper").drop_vars("uniform_endpoint")
    # get the cases where jumps occur
    pit_jumps = upper_values.where(upper_values == lower_values)

    cdf_left = xr.zeros_like(pit_jumps).where(x_values <= pit_jumps, 1).where(pit_jumps.notnull())
    cdf_right = xr.zeros_like(pit_jumps).where(x_values < pit_jumps, 1).where(pit_jumps.notnull())
    return {"left": cdf_left, "right": cdf_right}


def _pit_distribution_for_unif(pit_values: XarrayLike, x_values: xr.DataArray) -> XarrayLike:
    """
    Gives the values F(x) where F is the CDF for the PIT of a particular forcast--observation case,
    and x comes from `x_values`, given that the CDF is uniform. This occurs precisely
    when upper > lower in the [lower, upper] representation of the PIT value.
    If this condition fails, NaNs are returned.

    Left-hand and right-hand limits of F(x) are equal in this case.

    It is assumed that `x_values` containes all the values in `pit_values`.

    Args:
        pit_values: output from `_pit_values_for_ens` or `_pit_values_for_cdf`, containing
            dimension 'uniform_endpoint' with coordinates ['lower', 'upper']
        x_values: xr.DataArray of x values, with dimension 'pit_x_value', output from
            `_get_pit_x_values`

    Returns:
        An xarray object representing values of F(x), where x is represented by values
        in the 'pit_x_value' dimension.
    """
    lower_values = pit_values.sel(uniform_endpoint="lower").drop_vars("uniform_endpoint")
    upper_values = pit_values.sel(uniform_endpoint="upper").drop_vars("uniform_endpoint")
    # get the cases where the cdf is a uniform distribution over [a,b], a < b
    unif_cases = upper_values > lower_values
    # get the upper and lower values: these correspond to a, b
    pit_unif_upper = upper_values.where(unif_cases)  # the b
    pit_unif_lower = lower_values.where(unif_cases)  # the a
    # calculate the cdf of Unif(a,b) at the values in x
    pit_unif = (
        xr.zeros_like(pit_unif_upper)
        .where(x_values <= pit_unif_lower)  # 0 for x <= a
        .where(x_values < pit_unif_upper, 1)  # 1 for x >= b
        .interpolate_na("pit_x_value")  # interpolate in between
        .where(unif_cases)  # preserve nans where appropriate
    )
    return pit_unif


def _pit_cdfvalues(pit_values: XarrayLike) -> dict:
    """
    Calculates F(x) for each CDF F representing the PIT value for a forecast
    case. The x values used are all values where there is a non-linear change
    at least one F from the collection of Fs. Intermediate values of F can be recovered
    using linear interpolation.

    The values are output as a dictionary of two xarray object representing left hand
    and right hand limits of F at the point x.

    Args:
        pit_values: output from `_pit_values_for_ens` or `_pit_values_for_cdf`, containing
            dimension 'uniform_endpoint' with coordinates ['lower', 'upper']

    Returns:
        dictionary of cdf values, with keys 'left' and 'right',
        representing the left and right hand limits of the cdf values F(x).
        Each value in the dictionary is an xarray object representing limits of F(x),
        where x is represented by values in the 'pit_x_value' dimension.
    """
    # get the x values
    x_values = _get_pit_x_values(pit_values)

    # get cdf values where the pit cdf jumps
    cdf_at_jump_cases = _pit_distribution_for_jumps(pit_values, x_values)
    # get the cdf values where the pit is uniform
    cdf_at_unif_cases = _pit_distribution_for_unif(pit_values, x_values)
    # combine
    cdf_right = cdf_at_jump_cases["right"].combine_first(cdf_at_unif_cases)
    cdf_left = cdf_at_jump_cases["left"].combine_first(cdf_at_unif_cases)

    return {"left": cdf_left, "right": cdf_right}


def _pit_distribution_for_ens(
    fcst: XarrayLike,
    obs: XarrayLike,
    ens_member_dim: str,
    *,  # Force keywords arguments to be keyword-only
    reduce_dims: Optional[FlexibleDimensionTypes] = None,
    preserve_dims: Optional[FlexibleDimensionTypes] = None,
    weights: Optional[XarrayLike] = None,
) -> dict:
    """
    Each forecast case (i.e., an ensemble of forecasts) is interpreted as an empirical
    cumulative distribution function (CDF) :math:`G`. Given any observation :math:`y`,
    the probability integral transform (PIT) value at :math:`G` is a uniform distribution
    on the closed interval :math:`[G(y-), G(y)]`, where :math:`G(y-)` denotes the left-hand
    limit of :math:`G` at :math:`y`.

    This function outputs values of PIT by representing each uniform distribution on
    :math:`[G(y-), G(y)]` as the corresponding uniform CDF :math:`F`. We call :math:`F`
    the 'PIT CDF' for the forecast--observation pair :math:`(G,y)`. Values :math:`F(x-)`
    and :math:`F(x)` are given for an optimal set of points :math:`x` satisfying
    :math:`0 <= x <= 1`. The set is optimal in the sense that the value of :math:`F`
    elsewhere can be determined via linear interpolation, whilst no smaller set of values
    has this property.

    Dimensions are be reduced by taking (possibly weighted) means of the CDF values
    :math:`F(x)`. The weighted means can still be interpreted a values of CDFs, and
    intermediate values still attained via linear interpolation.

    Args:
        fcst: an xarray object of ensemble forecasts, containing the dimension
            `ens_member_dim`.
        obs: an xarray object of observations.
        ens_member_dim: name of the ensemble member dimension in ``fcst``.
        reduce_dims: Optionally specify which dimensions to reduce when calculating the
            PIT CDF values, where the mean is taken over all forecast cases.
            All other dimensions will be preserved. As a special case, 'all' will allow
            all dimensions to be reduced. Only one of ``reduce_dims`` and ``preserve_dims``
            can be supplied. The default behaviour if neither are supplied is to reduce all dims.
        preserve_dims: Optionally specify which dimensions to preserve when calculating the
            PIT CDF values, where the mean is taken over all forecast cases.
            All other dimensions will be reduced. As a special case, 'all' will allow
            all dimensions to be preserved, apart from ``severity_dim`` and ``prob_threshold_dim``.
            Only one of ``reduce_dims`` and ``preserve_dims`` can be supplied. The default
            behaviour if neither are supplied is to reduce all dims.
        weights: Optionally provide an array for weighted averaging (e.g. by area, by latitude,
            by population, custom) of PIT CDF values across all forecast cases.

    Returns:
        a dictionary with the following keys and values:
        - "left": an xarray object containing the left-hand limits :math:`F(x-)` of the PIT CDF values
        - "right": an xarray object containing the values :math:`F(x)` of the PIT CDF values
        - "pit_uniform_endpoints": an xarray object containing the values :math:`[G(y-), G(y)]`,
            with lower and upper endpoints indexed by the dimension "uniform_endpoint"

    Raises:
        ValueError if dimensions of ``fcst``, ``obs`` or ``weights`` contain any of the following reserved names:
                'uniform_endpoint', 'pit_x_value', 'x_plotting_position', 'y_plotting_position', 'plotting_point'
    """

    dims_for_mean = _dims_for_mean_with_checks(fcst, obs, ens_member_dim, weights, reduce_dims, preserve_dims)

    # PIT values in [G(y-), G(y)] format
    pit_values = _pit_values_for_ens(fcst, obs, ens_member_dim)
    # convert to CDF format, take weighted means, and output dictionary
    result = _pit_values_final_processing(pit_values, weights, dims_for_mean)

    return result


################################################
# Functions for Pit and PitFcstAtObs methods
################################################


def _get_plotting_points(left: XarrayLike, right: XarrayLike) -> XarrayLike:
    """
    Given left- and right-hand limiting values of the PIT CDF
    (e.g. as output by `Pit().left` and `Pit().right`), concatenates these objects along
    the "pit_x_value" and sorts them along that dimension.
    The "pit_x_value" index will have duplicate values.
    """
    return xr.concat([left, right], "pit_x_value").sortby("pit_x_value")


def _get_plotting_points_param(left: XarrayLike, right: XarrayLike) -> dict:
    """
    Given outputs left and right from `pit_cdfvalues`, calculates the plotting positions
    for the PIT CDF (equivalently PIT-uniform probability plot).

    Args:
        left: "left" output from `pit_cdfvalues`, namely the left-hand limits of the PIT CDF
        right: "right" output from `pit_cdfvalues`, namely values of the PIT CDF

    Returns:
        dictionary with following keys and values:
            'x_plotting_position': xarray object with the x-axis plotting positions for each
                point, indexed by 'plotting_point'
            'y_plotting_position': xarray object with the y-axis plotting positions for each
                point, indexed by 'plotting_point'
    """
    # only keep coordinates in left where (left != right) somewhere and where left is not NaN.
    dims_for_any = set(left.dims) - {"pit_x_value"}
    different = ((left != right) & left.notnull()).any(dims_for_any)
    left_reduced = left.where(different).dropna("pit_x_value", how="all")

    y_points = _get_plotting_points(left_reduced, right)

    if isinstance(y_points, xr.Dataset):
        x_points = xr.merge([y_points[var]["pit_x_value"].rename(var) for var in y_points.data_vars])
    else:
        x_points = y_points["pit_x_value"]
    # reindex with a plotting point index
    y_points = y_points.assign_coords(pit_x_value=np.arange(len(y_points.pit_x_value))).rename(
        {"pit_x_value": "plotting_point"}
    )
    x_points = x_points.assign_coords(pit_x_value=np.arange(len(x_points.pit_x_value))).rename(
        {"pit_x_value": "plotting_point"}
    )
    return {"x_plotting_position": x_points, "y_plotting_position": y_points}


def _value_at_pit_cdf(pit_left: XarrayLike, pit_right: XarrayLike, point: float) -> XarrayLike:
    """
    Given mean PIT CDF `F`, with left-hand limit values `pit_left` and right-hand limit values
    `pit_right`, finds the value `F(point)` under the assumption that `point` is not in
    `pit_left["pit_x_value"]`.

    The result is attained via linear interpolation between the known values at the
    nearest right-hand limit (less than `point`) and the nearest left-hand limit
    (greater than `point`).

    Args:
        pit_left: `left` attribute of a `Pit_for_ensemble` object
        pit_right: `right` attribute of a `Pit_for_ensemble` object
        point: a float strictly between 0 and 1

    Returns:
        xarray object containing values of the CDF `F` at `point`, with same dimensions
        as `pit_left`

    Raises:
        ValueError if `point` is in `pit_left["pit_x_value"].values`
    """
    if point in pit_left["pit_x_value"].values:
        raise ValueError('`point` must not be a value in `pit_left["pit_x_value"]`')
    upper = np.min([x for x in pit_left["pit_x_value"].values if x > point])
    lower = np.max([x for x in pit_right["pit_x_value"].values if x < point])
    # get the values at upper and lower
    value = xr.concat([pit_right.sel(pit_x_value=lower), pit_left.sel(pit_x_value=upper)], "pit_x_value")
    # value at newpoint is attained via linear interpolation
    value = add_thresholds(value, "pit_x_value", [point], "linear").sel(pit_x_value=slice(point, point))
    return value


def _construct_hist_values(cdf_at_endpoints: list[XarrayLike], bins: int) -> XarrayLike:
    """
    Calculates the PIT histogram values given values of the PIT CDF at the endpoint of
    every histogram bin.

    Args:
        cdf_at_endpoints: a list of xarray objects with dimension "pit_x_value",
            each giving the value of the PIT CDF at an endpoint of one of the histogram bins.
        bins: number of bins in the histogram.

    Returns:
        xarray object including dim 'bin_centre' and coordinates 'bin_left_endpoint',
        'bin_right_endpoint', with values the hight of each bar in the histogram.
    """
    # cdf values at all histogram bin endpoints
    cdf_at_endpoints = xr.concat(cdf_at_endpoints, "pit_x_value").sortby("pit_x_value")

    # calculate the histogram values
    histogram_values = cdf_at_endpoints.diff("pit_x_value")

    histogram_values = histogram_values.assign_coords(
        {
            "pit_x_value": histogram_values["pit_x_value"] - 1 / (2 * bins),
            "bin_left_endpoint": histogram_values["pit_x_value"] - 1 / bins,
            "bin_right_endpoint": histogram_values["pit_x_value"],
        }
    ).rename({"pit_x_value": "bin_centre"})
    return histogram_values


def _pit_hist_left(pit_left: XarrayLike, pit_right: XarrayLike, bins: int) -> XarrayLike:
    """
    Gets the pit histogram values, with left endpoints included in every bin.
    e.g. if bins=5, the bins are [0, 0.2), [0.2, 0.4), ..., [0.8, 1]

    Args:
        pit_left: `left` attribute of a `Pit_for_ensemble` object
        pit_right: `right` attribute of a `Pit_for_ensemble` object
        bins: the number of bins in the histogram

    Returns:
        xarray object including dim 'bin_centre' and coordinates 'bin_left_endpoint',
        'bin_right_endpoint', with values the hight of each bar in the histogram.
    """
    left_endpoints = np.arange(bins) / bins

    # the value of the mean PIT CDF at 1 is 1
    # using .sel is safe because the value 1.0 is guaranteed to be in pit_x_value (see _get_pit_x_values)
    value_at_1 = pit_right.sel(pit_x_value=slice(1.0, 1.0))

    # start collecting the value of the pit CDF at each of the endpoints of the histogram bins
    cdf_at_endpoints = [value_at_1]

    # get values of the PIT CDF that are aleady available in pit
    common_points = [x for x in left_endpoints if x in pit_left["pit_x_value"]]
    # common_points is no-empty because 0.0 is in both left_endpoints and pit_left["pit_x_value"];
    # see _get_pit_x_values
    cdf_at_common_points = pit_left.sel(pit_x_value=common_points)
    cdf_at_endpoints.append(cdf_at_common_points)

    # get values of the PIT CDF that are not available in pit
    new_points = [x for x in left_endpoints if x not in pit_left["pit_x_value"]]
    for new_point in new_points:
        value = _value_at_pit_cdf(pit_left, pit_right, new_point)
        cdf_at_endpoints.append(value)

    histogram_values = _construct_hist_values(cdf_at_endpoints, bins)
    return histogram_values


def _pit_hist_right(pit_left: XarrayLike, pit_right: XarrayLike, bins: int) -> XarrayLike:
    """
    Gets the pit histogram values, with right endpoints included in every bin.
    e.g. if bins=5, the bins are [0, 0.2], (0.2, 0.4], ..., (0.8, 1]

    Args:
        pit_left: `left` attribute of a `Pit_for_ensemble` object
        pit_right: `right` attribute of a `Pit_for_ensemble` object
        bins: the number of bins in the histogram

    Returns:
        xarray object including dim 'bin_centre' and coordinates 'bin_left_endpoint',
        'bin_right_endpoint', with values the hight of each bar in the histogram.
    """
    right_endpoints = np.arange(1, bins + 1) / bins

    # the value of the mean PIT CDF at 0 is 0, except where NaNs
    # using .sel is safe because the value 0.0 is guaranteed to be in pit_x_value (see _get_pit_x_values)
    value_at_0 = pit_left.sel(pit_x_value=slice(0.0, 0.0))

    cdf_at_endpoints = [value_at_0]

    # get values of the PIT CDF that are aleady available in pit
    common_points = [x for x in right_endpoints if x in pit_right["pit_x_value"]]
    # common_points is non-empty because 1.0 is in both right_endpoints and pit_left["pit_x_value"];
    # see _get_pit_x_values and floating point arithmetic tested on
    # (np.arange(1000) / np.arange(1000)) == 1.
    cdf_at_common_points = pit_right.sel(pit_x_value=common_points)
    cdf_at_endpoints.append(cdf_at_common_points)

    # get values of the PIT CDF that are not available in pit
    new_points = [x for x in right_endpoints if x not in pit_right["pit_x_value"]]
    for new_point in new_points:
        value = _value_at_pit_cdf(pit_left, pit_right, new_point)
        cdf_at_endpoints.append(value)

    histogram_values = _construct_hist_values(cdf_at_endpoints, bins)

    return histogram_values


def _alpha_score(left: XarrayLike, right: XarrayLike, negative_orientation: bool = True) -> XarrayLike:
    """
    Given left-hand and right-hand limits of the Pit CDF (.left and .right attributes),
    calculates the alpha score.

    Args:
        left: xarray object of left-hand limits with 'pit_x_value' dimension.
        right: xarray object of right-hand limits with 'pit_x_value' dimension.

    Returns:
        xarray object with alpha score, and 'pit_x_value' dimension collapsed
    """
    if isinstance(left, xr.DataArray):
        return _alpha_score_array(left, right, negative_orientation=negative_orientation)

    return xr.merge(
        [
            _alpha_score_array(left[var], right[var], negative_orientation=negative_orientation).rename(var)
            for var in left.data_vars
        ]
    )


def _alpha_score_array(left: xr.DataArray, right: xr.DataArray, negative_orientation: bool = True) -> xr.DataArray:
    """
    The alpha score with negative orientation is integrated absolute difference between
    the graph of the PIT CDF and the diagonal. This will be calculated using the trapezoidal rule.
    For the calculation to be exact, we need to include points where the PIT CDF and diagonal intersect.

    This function does this when the left and right limits of the PIT CDF are data arrays.

    When negative_orientation=False, the alpha score is 1 - twice the integrated difference.

    Args:
        left: left-hand limit of the PIT CDF, indexed by "pit_x_value". Other dimensions possible.
        left: right-hand limit of the PIT CDF, indexed by "pit_x_value". Other dimensions possible.

    Returns:
        array of alpha scores, with "pit_x_value" dimension collapsed but other dimensions preserved.
    """
    # need to expand plotting_points to include points where the PIT CDF crosses the diagonal.
    # To do this, use interp1d since xarray via pandas interpolate cannot handle duplicate values in an index
    param_plotting_points = _get_plotting_points_param(left, right)
    intersection_points = _diagonal_intersection_points(param_plotting_points)
    plotting_points = _get_plotting_points(left, right)

    x_axis_num = plotting_points.get_axis_num("pit_x_value")
    y_values = plotting_points.values
    x_values = plotting_points["pit_x_value"].values

    plotting_points_interpolated = interp1d(x_values, y_values, axis=x_axis_num, kind="linear")(intersection_points)

    # turn the numpy array into an xarray array
    interpolated_coords = dict(plotting_points.coords)
    interpolated_coords["pit_x_value"] = intersection_points
    plotting_points_interpolated = xr.DataArray(
        data=plotting_points_interpolated, dims=plotting_points.dims, coords=interpolated_coords
    )
    plotting_points = xr.concat([plotting_points, plotting_points_interpolated], "pit_x_value").sortby("pit_x_value")

    score = np.abs(plotting_points - plotting_points["pit_x_value"]).integrate("pit_x_value")

    if negative_orientation:
        return score

    return 1 - 2 * score


def _diagonal_intersection_points(param_plotting_points: dict) -> np.ndarray:
    """
    Gets the x values where the line y = x intersects with the piecewise linear graph
    y = F(x), where F is the PIT CDF. If the graph of F is discontinuous (vertical)
    at a point of intersection, or if F is coincident with the diagonal over an open interval,
    then the corresponding points of intersection are not returned as they are not needed to
    calculate the alpha score.

    Only handles param_plotting_points in xarray data array format.

    Args:
        param_plotting_points: diction consisting of xarray data array output from
            `_get_plotting_points_param`

    Returns:
        1-dimensions numpy array of intersection points
    """
    x_pos = param_plotting_points["x_plotting_position"]
    y_pos = param_plotting_points["y_plotting_position"]

    diff_x = x_pos.diff("plotting_point")
    diff_y = y_pos.diff("plotting_point")
    nonzero_dx = diff_x != 0
    safe_diff_x = xr.where(nonzero_dx, diff_x, 1)

    # gradient of chord AB where A(x_pos[i-1], y_pos[i-1]), B(x_pos[i], y_pos[i])
    gradient = xr.where(nonzero_dx, diff_y / safe_diff_x, np.nan)
    # solution if there is a desired point of intersection, obtained by solving
    # simultaneous equations for equation of line AB with line x = y
    denom = 1 - gradient
    nonzero_denom = denom != 0
    safe_denom = xr.where(nonzero_denom, denom, 1)
    x_solution = xr.where(nonzero_dx & nonzero_denom, (y_pos - gradient * x_pos) / safe_denom, np.nan)
    # for x_solution to be of interest, require that  x_pos[i-1] < x_solution[i] < x_pos[i]
    x_solution = x_solution.where((x_solution < x_pos) & (x_solution > x_pos.shift(plotting_point=1)))
    x_solution = np.unique(x_solution.values.flatten())
    return x_solution[~np.isnan(x_solution)]


def _expected_value(plotting_points: XarrayLike) -> XarrayLike:
    """
    Calculates the expected value of the PIT distribution, given plotting points for the
    CDF of that distribution. Uses the well-known formula that the expected value of a
    non-negative random variable is the integral of 1 - CDF.

    Here, `plotting_points`, with linear interpolation, fully describes the CDF of the
    PIT distribution.

    Args:
        plotting_points: xarray object of plotting points, indexed by non-decreasing values
            (with duplicates) along the 'pit_x_value' dimension.

    Returns:
        xarray object with expected value, and 'pit_x_value' dimension collapsed
    """
    return (1 - plotting_points).integrate("pit_x_value")


def _variance(plotting_points: XarrayLike) -> XarrayLike:
    """
    Calculates the variance of the PIT distribution using the formula
        Var(Y) = 2 * integral(x * (1 - F(x)), x >= 0) + E(Y)^2,
    for a non-negative random variable Y, where F is the CDF of Y.

    Here, `plotting_points`, with linear interpolation, fully describes the CDF of the
    PIT distribution.

    Args:
        plotting_points: xarray object of plotting points, indexed by non-decreasing values
            (with duplicates) along the 'pit_x_value' dimension.

    Returns:
        xarray object with the variance, and 'pit_x_value' dimension collapsed
    """
    expected_value = _expected_value(plotting_points)
    integral_term = 2 * _variance_integral_term(plotting_points)
    return integral_term - expected_value**2


def _variance_integral_term(plotting_points: XarrayLike) -> XarrayLike:
    """
    Calculates
        integral(x * (1 - F(x)), x >= 0)
    where F is a CDF.

    Here, `plotting_points`, with linear interpolation, fully describes the CDF of the
    PIT distribution F. It is assumed coordinates in `plotting_points['pit_x_value']`
    come in duplicate pairs (e.g. [0, 0, 0.1, 0.1, 0.15, 0.15, ...]), to describe
    left-hand and right-hand limits at each point. This assumption holds with output
    from the `.plotting_points` method.

    Args:
        plotting_points: xarray object of plotting points, indexed by non-decreasing values
            (with duplicates) along the 'pit_x_value' dimension.

    Returns:
        xarray object with the integral values, and 'pit_x_value' dimension collapsed
    """
    # notation: F_i(t) = m_i * t + b_i whenever x[i-1] <= t <= x[i]
    # and x[i] is a value in `function_values[threshold_dim]`.

    x_values = plotting_points["pit_x_value"]
    x_shifted = plotting_points["pit_x_value"].shift(pit_x_value=1)
    # difference in x values
    diff_xs = x_values - x_shifted
    # difference in function values y_i = F(x[i])
    diff_ys = plotting_points - plotting_points.shift(pit_x_value=1)

    nonzero = diff_xs != 0
    safe_diff_xs = xr.where(nonzero, diff_xs, np.nan)

    # gradients m
    m_values = xr.where(nonzero, diff_ys / safe_diff_xs, 0)
    # intercepts b_i
    b_values = xr.where(nonzero, plotting_points - m_values * x_values, 0)
    # integral(t * (1 - F(t))) on the interval (x[i-1], x[i]), for each i, using calculus:
    integral_i = xr.where(
        nonzero,
        (1 - b_values) * (x_values**2 - x_shifted**2) / 2 - m_values * (x_values**3 - x_shifted**3) / 3,
        0,
    )

    integral = integral_i.sum("pit_x_value")
    # return NaN if NaN in function_values
    integral = integral.where(~np.isnan(plotting_points).any("pit_x_value"))

    return integral
