text: #   Copyright 2020 The PyMC Developers
  #
  #   Licensed under the Apache License, Version 2.0 (the "License");
  #   you may not use this file except in compliance with the License.
  #   You may obtain a copy of the License at
  #
  #       http://www.apache.org/licenses/LICENSE-2.0
  #
  #   Unless required by applicable law or agreed to in writing, software
  #   distributed under the License is distributed on an "AS IS" BASIS,
  #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #   See the License for the specific language governing permissions and
  #   limitations under the License.
  import aesara.tensor as at
  import numpy as np
  
  from aesara.scalar import Clip
  from aesara.tensor import TensorVariable
  from aesara.tensor.random.op import RandomVariable
  
  from pymc.aesaraf import change_rv_size
  from pymc.distributions.distribution import SymbolicDistribution, _moment
  from pymc.util import check_dist_not_registered
  
  
  class Censored(SymbolicDistribution):
      r"""
      Censored distribution
  
      The pdf of a censored distribution is
  
      .. math::
  
          \begin{cases}
              0 & \text{for } x < lower, \\
              \text{CDF}(lower, dist) & \text{for } x = lower, \\
              \text{PDF}(x, dist) & \text{for } lower < x < upper, \\
              1-\text{CDF}(upper, dist) & \text {for} x = upper, \\
              0 & \text{for } x > upper,
          \end{cases}
  
  
      Parameters
      ----------
      dist: unnamed distribution
          Univariate distribution created via the `.dist()` API, which will be censored.
          This distribution must have a logcdf method implemented for sampling.
  
          .. warning:: dist will be cloned, rendering it independent of the one passed as input.
  
      lower: float or None
          Lower (left) censoring point. If `None` the distribution will not be left censored
      upper: float or None
          Upper (right) censoring point. If `None`, the distribution will not be right censored.
  
      Warnings
      --------
      Continuous censored distributions should only be used as likelihoods.
      Continuous censored distributions are a form of discrete-continuous mixture
      and as such cannot be sampled properly without a custom step sampler.
      If you wish to sample such a distribution, you can add the latent uncensored
      distribution to the model and then wrap it in a :class:`~pymc.Deterministic`
      :func:`~pymc.math.clip`.
  
  
      Examples
      --------
      .. code-block:: python
  
          with pm.Model():
              normal_dist = pm.Normal.dist(mu=0.0, sigma=1.0)
              censored_normal = pm.Censored("censored_normal", normal_dist, lower=-1, upper=1)
      """
  
      @classmethod
      def dist(cls, dist, lower, upper, **kwargs):
          if not isinstance(dist, TensorVariable) or not isinstance(dist.owner.op, RandomVariable):
              raise ValueError(
                  f"Censoring dist must be a distribution created via the `.dist()` API, got {type(dist)}"
              )
          if dist.owner.op.ndim_supp > 0:
              raise NotImplementedError(
                  "Censoring of multivariate distributions has not been implemented yet"
              )
          check_dist_not_registered(dist)
          return super().dist([dist, lower, upper], **kwargs)
  
      @classmethod
      def num_rngs(cls, *args, **kwargs):
          return 1
  
      @classmethod
      def ndim_supp(cls, *dist_params):
          return 0
  
      @classmethod
      def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None):
  
          lower = at.constant(-np.inf) if lower is None else at.as_tensor_variable(lower)
          upper = at.constant(np.inf) if upper is None else at.as_tensor_variable(upper)
  
          # When size is not specified, dist may have to be broadcasted according to lower/upper
          dist_shape = size if size is not None else at.broadcast_shape(dist, lower, upper)
          dist = change_rv_size(dist, dist_shape)
  
          # Censoring is achieved by clipping the base distribution between lower and upper
          rv_out = at.clip(dist, lower, upper)
  
          # Reference nodes to facilitate identification in other classmethods, without
          # worring about possible dimshuffles
          rv_out.tag.dist = dist
          rv_out.tag.lower = lower
          rv_out.tag.upper = upper
  
          if rngs is not None:
              rv_out = cls._change_rngs(rv_out, rngs)
  
          return rv_out
  
      @classmethod
      def _change_rngs(cls, rv, new_rngs):
          (new_rng,) = new_rngs
          dist_node = rv.tag.dist.owner
          lower = rv.tag.lower
          upper = rv.tag.upper
          olg_rng, size, dtype, *dist_params = dist_node.inputs
          new_dist = dist_node.op.make_node(new_rng, size, dtype, *dist_params).default_output()
          return cls.rv_op(new_dist, lower, upper)
  
      @classmethod
      def change_size(cls, rv, new_size, expand=False):
          dist = rv.tag.dist
          lower = rv.tag.lower
          upper = rv.tag.upper
          new_dist = change_rv_size(dist, new_size, expand=expand)
          return cls.rv_op(new_dist, lower, upper)
  
  
  @_moment.register(Clip)
  def moment_censored(op, rv, dist, lower, upper):
      moment = at.switch(
          at.eq(lower, -np.inf),
          at.switch(
              at.isinf(upper),
              # lower = -inf, upper = inf
              0,
              # lower = -inf, upper = x
              upper - 1,
          ),
          at.switch(
              at.eq(upper, np.inf),
              # lower = x, upper = inf
              lower + 1,
              # lower = x, upper = x
              (lower + upper) / 2,
          ),
      )
      moment = at.full_like(dist, moment)
      return moment
  