text: from abc import ABC, abstractmethod
  from functools import wraps
  from typing import Tuple
  
  import numpy as np
  from scipy.special import ndtr, ndtri
  from scipy.stats._continuous_distns import _norm_pdf
  
  from ..mapper.operator import (
      DiagonalMatrix,
      LinearOperator,
      ShermanMorrison
  )
  
  
  def numerical_jacobian(x, func, eps=1e-8, args=(), **kwargs):
      x = np.array(x)
      f0 = func(x, *args, **kwargs)
      jac = np.empty(np.shape(f0) + np.shape(x))
      fslice = (slice(None),) * np.ndim(f0)
      with np.nditer(x, flags=['multi_index'], op_flags=['readwrite']) as it:
          for xi in it:
              xi += eps
              f1 = func(x, *args, **kwargs)
              jac[fslice + it.multi_index] = (f1 - f0) / eps
              xi -= eps
  
      return jac
  
  
  class AbstractDensityTransform(ABC):
      """
      This class allows the transformation of a probability density function, p(x)
      whilst preserving the measure of the distribution, i.e.
  
      \int p(x) dx = 1
  
      p'(f) = p(f(x)) * |df/dx|
  
      \inf p'(f) df = 1
  
      Methods
      -------
      transform
          calculates f(x)
      
      inv_transform
          calculates f^{-1}(y)
  
      jacobian
          calculates df/dx
  
      log_det
          calculates log |df/dx|
  
      log_det_grad
          calculates |df/dx|, d log_det/dx
  
      transform_det
          calculates f(x), |df/dx|
  
      transform_jac
          calculates f(x), df/dx
  
      transform_det_jac
          calculates f(x), log_det, d log_det/dx, df/dx
  
      These final 3 functions are defined so that child classes
      can define custom methods that avoid recalculation of intermediate 
      values that are needed to calculate multiple versions of the quantities
      """
  
      @abstractmethod
      def transform(self, x):
          pass
  
      @abstractmethod
      def inv_transform(self, x):
          pass
  
      @abstractmethod
      def jacobian(self, x: np.ndarray) -> LinearOperator:
          pass
  
      @abstractmethod
      def log_det(self, x: np.ndarray) -> np.ndarray:
          pass
  
      @abstractmethod
      def log_det_grad(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
          pass
  
      def _numerical_jacobian_func(self, x):
          return self.log_det_grad(x)[1].sum(0)
  
      def log_det_hess(self, x: np.ndarray) -> np.ndarray:
          return numerical_jacobian(
              x, self._numerical_jacobian_func
          )
  
      def transform_det(self, x) -> Tuple[np.ndarray, np.ndarray]:
          return self.transform(x), self.log_det(x)
  
      def transform_jac(self, x) -> Tuple[np.ndarray, LinearOperator]:
          return self.transform(x), self.jacobian(x)
  
      def transform_det_jac(
              self, x
      ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, LinearOperator]:
          return (
              self.transform(x),
              *self.log_det_grad(x),
              self.jacobian(x)
          )
  
      def transform_func(self, func):
          @wraps(func)
          def transformed_func(*args, **kwargs):
              x, *args = args
              x = self.transform(x)
              return func(x, *args, **kwargs)
  
          transformed_func.transform = self
  
          return transformed_func
  
      def transform_func_grad(self, func_grad):
          @wraps(func_grad)
          def transformed_func_grad(*args, **kwargs):
              x, *args = args
              x, jac = self.transform_jac(x)
              val, grad = func_grad(x, *args, **kwargs)
              return x, grad * jac
  
          transformed_func_grad.transform = self
          return transformed_func_grad
  
      def transform_func_grad_hess(self, func_grad_hess):
          @wraps(func_grad_hess)
          def transformed_func_grad_hess(*args, **kwargs):
              x, *args = args
              x, jac = self.transform_jac(x)
              val, grad, hess = func_grad_hess(x, *args, **kwargs)
              return val, grad * jac, jac.quad(hess)
  
          transformed_func_grad_hess.transform = self
          return transformed_func_grad_hess
  
  
  class LinearTransform(AbstractDensityTransform):
      def __init__(self, linear: LinearOperator):
          self.linear = linear
  
      def transform(self, x: np.ndarray) -> np.ndarray:
          return self.linear * x
  
      def inv_transform(self, x: np.ndarray) -> np.ndarray:
          return self.linear.ldiv(x)
  
      def jacobian(self, x: np.ndarray) -> np.ndarray:
          return self.linear
  
      def log_det(self, x: np.ndarray) -> np.ndarray:
          return self.linear.log_det
  
      def log_det_grad(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
          return self.log_det(x), 0
  
  
  class LinearShiftTransform(LinearTransform):
      def __init__(self, shift: float = 0, scale: float = 1):
          self.shift = shift
          self.scale = scale
          super().__init__(
              DiagonalMatrix(np.reciprocal(self.scale))
          )
  
      def inv_transform(self, x: np.ndarray) -> np.ndarray:
          return x * self.scale + self.shift
  
      def transform(self, x: np.ndarray) -> np.ndarray:
          return (x - self.shift) / self.scale
  
      def log_det(self, x: np.ndarray) -> np.ndarray:
          return - np.log(self.scale) * np.ones_like(x)
  
  
  class FunctionTransform(AbstractDensityTransform):
      def __init__(self, func, inv_func, grad, hess=None, args=(), func_grad_hess=None):
          self.func = func
          self.inv_func = inv_func
          self.grad = grad
          self.hess = hess
          self.args = args
          self.func_grad_hess = func_grad_hess
  
      def transform(self, x):
          return self.func(x, *self.args)
  
      def inv_transform(self, x):
          return self.inv_func(x, *self.args)
  
      def jacobian(self, x):
          return DiagonalMatrix(self.grad(x, *self.args))
  
      def log_det(self, x: np.ndarray) -> np.ndarray:
          gs = self.grad(x, *self.args)
          return np.log(gs)
  
      def log_det_grad(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
          if self.func_grad_hess:
              x0, gs, hs = self.func_grad_hess(x, *self.args)
          else:
              gs = self.grad(x, *self.args)
              hs = self.hess(x, *self.args)
          return np.log(gs), hs / gs
  
      def transform_det_jac(
              self, x
      ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, LinearOperator]:
          if self.func_grad_hess:
              x0, gs, hs = self.func_grad_hess(x, *self.args)
          else:
              x0 = self.func(x, *self.args)
              gs = self.grad(x, *self.args)
              hs = self.hess(x, *self.args)
          return x0, np.log(gs), hs / gs, DiagonalMatrix(gs)
  
  
  def exp3(x):
      expx = np.exp(x)
      return (expx,) * 3
  
  
  exp_transform = FunctionTransform(np.exp, np.log, np.exp, func_grad_hess=exp3)
  
  
  def log3(x):
      ix = np.reciprocal(x)
      return np.log(x), ix, -np.square(ix)
  
  
  log_transform = FunctionTransform(
      np.log, np.exp, np.reciprocal, func_grad_hess=log3)
  
  
  def _log_10_inv(x):
      return 10 ** x
  
  
  # TODO: what should func_grad_hess look like? np.log10, ... ?
  log_10_transform = FunctionTransform(
      np.log10, _log_10_inv, np.reciprocal
  )
  
  
  def sigmoid(x, scale=1, shift=0):
      return scale / (1 + np.exp(-x)) + shift
  
  
  def logit(x, scale=1, shift=0):
      x = (x - shift) / scale
      return np.log(x) - np.log1p(-x)
  
  
  def sigmoid_grad(x, scale=1, shift=0):
      expx = np.exp(-x)
      return scale * expx / np.square(1 + expx)
  
  
  def logit_grad(x, scale=1, shift=0):
      x = (x - shift) / scale
      return (np.reciprocal(x) + np.reciprocal(1 - x)) / scale
  
  
  def logit_hess(x, scale=1, shift=0):
      x = (x - shift) / scale
      return np.reciprocal(1 - x) - np.reciprocal(x)
  
  
  def logit_grad_hess(x, scale=1, shift=0):
      x = (x - shift) / scale
      ix = np.reciprocal(x)
      ix1 = np.reciprocal(1 - x)
      ix2 = np.square(ix)
      ix12 = np.square(ix1)
      return (np.log(x) - np.log1p(-x), (ix + ix1) / scale, (ix12 - ix2) / scale ** 2)
  
  
  logistic_transform = FunctionTransform(
      logit, sigmoid, logit_grad, func_grad_hess=logit_grad_hess)
  
  
  def shifted_logistic(shift=0, scale=1):
      return FunctionTransform(logit, sigmoid, logit_grad, func_grad_hess=logit_grad_hess, args=(scale, shift))
  
  
  def ndtri_grad(x):
      return np.reciprocal(_norm_pdf(ndtri(x)))
  
  
  def ndtri_grad_hess(x):
      f = ndtri(x)
      phi = _norm_pdf(f)
      grad = np.reciprocal(phi)
      hess = grad ** 2 * f
      return f, grad, hess
  
  
  phi_transform = FunctionTransform(
      ndtri, ndtr, ndtri_grad, func_grad_hess=ndtri_grad_hess)
  
  
  class MultinomialLogitTransform(AbstractDensityTransform):
      """
      makes multinomial logististic transform from the p to x, where,
  
      x_i = log(p_i / (1 - sum(p)))
      p_i = exp(x_i) / (sum(exp(x_j) for x_j in x) + 1)
  
      When p's n-simplex is defined by,
  
      all(0 <= p_i <= 1 for p_i in p) and sum(p) < 1 
      """
  
      def __init__(self, axis=-1):
          self.axis = axis
  
      def _validate(self, p):
          p = np.asanyarray(p)
          keepdims = np.ndim(p) == self.ndim + 1
          if not (keepdims or np.ndim(p) == self.ndim):
              raise ValueError(
                  f"dimension of input must be {self.ndim} or {self.ndim + 1}")
  
          return p, keepdims
  
      def transform(self, p):
          p = np.asanyarray(p)
          lnp1 = np.log(1 - np.sum(p, axis=self.axis, keepdims=True))
          lnp = np.log(p)
          return lnp - lnp1
  
      def inv_transform(self, x):
          expx = np.exp(x)
          return expx / (expx.sum(axis=self.axis, keepdims=True) + 1)
  
      def jacobian(self, p):
          p = np.asanyarray(p)
          pn1 = 1 - np.sum(p, axis=-1, keepdims=True)
          # ln1p = np.log(pn1)
          # lnp = np.log(p)
          return ShermanMorrison(
              DiagonalMatrix(1 / p),
              1 / np.sqrt(pn1) * np.ones_like(p)
          )
  
      def log_det(self, p):
          p = np.asanyarray(p)
          p1 = 1 - np.sum(p, axis=self.axis, keepdims=True)
          # Hack to make sure summation broadcasting works correctly
          log_d = (
                          - np.log(p).sum(axis=self.axis, keepdims=True) - np.log(p1)
                  ) * np.full_like(p, p1.size / p.size)
          return log_d
  
      def log_det_grad(self, p):
          p = np.asanyarray(p)
          p1 = 1 - np.sum(p, axis=self.axis, keepdims=True)
          # Hack to make sure summation broadcasting works correctly
          log_d = (
                          - np.log(p).sum(axis=self.axis, keepdims=True) - np.log(p1)
                  ) * np.full_like(p, p1.size / p.size)
          return log_d, 1 / p1 - 1 / p
  
      def transform_det_jac(
              self, p
      ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, LinearOperator]:
          p = np.asanyarray(p)
          pn1 = 1 - np.sum(p, axis=self.axis, keepdims=True)
          ln1p = np.log(pn1)
          lnp = np.log(p)
          x = lnp - ln1p
          # Hack to make sure summation broadcasting works correctly
          logd = (
                         - lnp.sum(axis=self.axis, keepdims=True) - ln1p
                 ) * np.full_like(p, pn1.size / p.size)
          logd_grad = 1 / pn1 - 1 / p
          jac = ShermanMorrison(
              DiagonalMatrix(1 / p),
              1 / np.sqrt(pn1) * np.ones_like(p)
          )
          return (
              x, logd, logd_grad, jac
          )
  
  
  multinomial_logit_transform = MultinomialLogitTransform()
  