from __future__ import division

import numpy as np
from utils import msqrt, check_type, round_array, float_dtypes, integer_dtypes, bool_dtypes, safe_len, find_generations, logp_of_set, symmetrize
from numpy import ones, zeros, log, shape, cov, ndarray, inner, reshape, sqrt, any, array, all, abs, exp, where, isscalar, iterable, multiply, transpose, tri
from numpy.linalg.linalg import LinAlgError
from numpy.linalg import pinv, cholesky
from numpy.random import randint, random
from numpy.random import normal as rnormal
from numpy.random import poisson as rpoisson
from PyMCObjects import Stochastic, Potential, Deterministic
from Container import Container
from Node import ZeroProbability, Node, Variable, StochasticBase
from pymc.decorators import prop
import distributions
from copy import copy
from InstantiationDecorators import deterministic
import pdb, warnings, sys
import inspect

__docformat__='reStructuredText'

# Changeset history
# 22/03/2007 -DH- Added a _state attribute containing the name of the attributes that make up the state of the step method, and a method to return that state in a dict. Added an id.
# TODO: Test cases for binary and discrete Metropolises.

conjugate_Gibbs_competence = 0
nonconjugate_Gibbs_competence = 0

class AdaptationError(ValueError): pass

__all__=['DiscreteMetropolis', 'Metropolis', 'PDMatrixMetropolis', 'StepMethod', 'assign_method',  'pick_best_methods', 'StepMethodRegistry', 'NoStepper', 'BinaryMetropolis', 'AdaptiveMetropolis','Gibbs','conjugate_Gibbs_competence', 'nonconjugate_Gibbs_competence', 'DrawFromPrior']

StepMethodRegistry = []

def pick_best_methods(stochastic):
    """
    Picks the StepMethods best suited to handle
    a stochastic variable.
    """

    # Keep track of most competent methohd
    max_competence = 0
    # Empty set of appropriate StepMethods
    best_candidates = set([])

    # Loop over StepMethodRegistry
    for method in StepMethodRegistry:

        # Parse method and its associated competence
        try:
            competence = method.competence(stochastic)
        except:
#             print '\n\tWarning, there was an error while step method %s assessed its competence \n \
# \tto handle stochastic %s. It is being excluded from consideration.\n' \
#                     %(method.__name__, stochastic)
            competence = 0

        # If better than current best method, promote it
        if competence > max_competence:
            best_candidates = set([method])
            max_competence = competence

        # If same competence, add it to the set of best methods
        elif competence == max_competence:
            best_candidates.add(method)

    if max_competence<=0:
        raise ValueError, 'Maximum competence reported for stochastic %s is <= 0... you may need to write a custom step method class.' % stochastic.__name__

    # print s.__name__ + ': ', best_candidates, ' ', max_competence
    return best_candidates

def assign_method(stochastic, scale=None):
    """
    Returns a step method instance to handle a
    variable. If several methods have the same competence,
    it picks one arbitrarily (using set.pop()).
    """

    # Retrieve set of best candidates
    best_candidates = pick_best_methods(stochastic)

    # Randomly grab and appropriate method
    method = best_candidates.pop()

    failure_header = """Failed attempting to automatically assign step method class %s
to stochastic variable %s. Try setting %s's competence method to return 0
and manually assigning it when appropriate. See the user guide.

Error message: """%(method.__name__, stochastic.__name__, method.__name__)

    try:
        if scale:
            out = method(stochastic, scale = scale)
        else:
            out = method(stochastic)
    except:
        a,b,c = sys.exc_info()
        raise a, failure_header + b.message, c
    return out

class StepMethodMeta(type):
    """
    Automatically registers new step methods if they can be automatically assigned:
    if their init method has one and only one required argument.
    """
    def __init__(cls, name, bases, dict):
        type.__init__(cls, name, bases, dict)
        args, varargs, varkw, defaults = inspect.getargspec(cls.__init__)
        auto_assignment_OK = False
        if len(args) == 2:
            auto_assignment_OK = True
        elif len(args)>2:
            if defaults is not None:
                if len(defaults) == len(args)-2:
                    auto_assignment_OK = True
        elif len(args) == 1 and varargs is not None:
            auto_assignment_OK = True

        if auto_assignment_OK:
            StepMethodRegistry.append(cls)

class StepMethod(object):
    """
    This object knows how to make Stochastics take single MCMC steps.
    Its step() method will be called by Model at every MCMC iteration.

    :Parameters:
          -variables : list, array or set
            Collection of PyMCObjects

          - verbose (optional) : integer
            Level of output verbosity: 0=none, 1=low, 2=medium, 3=high. Setting to none (Default) allows verbosity to be set by sampler.

    Externally-accessible attributes:
      stochastics:   The Stochastics over which self has jurisdiction which have observed = False.
      children: The combined children of all Variables over which self has jurisdiction.
      parents:  The combined parents of all Nodes over which self has jurisdiction, as a set.
      loglike:  The summed log-probability of self's children conditional on all of self's
                  Variables' current values. These will be recomputed only as necessary.
                  This descriptor should eventually be written in C.

    Externally accesible methods:
      sample(): A single MCMC step for all the Stochastics over which self has
        jurisdiction. Must be overridden in subclasses.
      tune(): Tunes proposal distribution widths for all self's Stochastics.
      competence(s): Examines Stochastic instance s and returns self's
        competence to handle it, on a scale of 0 to 3.

    To instantiate a StepMethod called S with jurisdiction over a
    sequence/set N of Nodes:

      >>> S = StepMethod(N)

    :SeeAlso: Metropolis, Sampler.
    """

    __metaclass__ = StepMethodMeta

    def __init__(self, variables, verbose=None, tally=False):
        # StepMethod initialization

        if not iterable(variables) or isinstance(variables, Node):
            variables = [variables]

        self.stochastics = set()
        self.children = set()
        self.parents = set()
        self.tally = tally

        self._state = []
        self._tuning_info = []
        self.verbose = verbose

        # File away the variables
        for variable in variables:
            # Sort.

            if isinstance(variable,Stochastic) and not variable.observed:
                self.stochastics.add(variable)

        if len(self.stochastics)==0:
            raise ValueError, 'No stochastics provided.'

        # Find children, no need to find parents; each variable takes care of those.
        for variable in variables:
            self.children |= variable.children
            for parent in variable.parents.itervalues():
                if isinstance(parent, Variable):
                    self.parents.add(parent)

        self.children = set([])
        self.parents = set([])
        for s in self.stochastics:
            self.children |= s.extended_children
            self.parents |= s.extended_parents

        # Remove own stochastics from children and parents.
        self.children -= self.stochastics
        self.parents -= self.stochastics

        # self.markov_blanket is a list, because we want self.stochastics to have the chance to
        # raise ZeroProbability exceptions before self.children.
        self.markov_blanket = list(self.stochastics)+list(self.children)

        # ID string for verbose feedback
        self._id = self.__class__.__name__ + '_' + '_'.join([s.__name__ for s in self.stochastics])

    def step(self):
        """
        Specifies single step of step method.
        Must be overridden in subclasses.
        """
        pass

    @staticmethod
    def competence(s):
        """
        This function is used by Sampler to determine which step method class
        should be used to handle stochastic variables.

        Return value should be a competence
        score from 0 to 3, assigned as follows:

        0:  I can't handle that variable.
        1:  I can handle that variable, but I'm a generalist and
            probably shouldn't be your top choice (Metropolis
            and friends fall into this category).
        2:  I'm designed for this type of situation, but I could be
            more specialized.
        3:  I was made for this situation, let me handle the variable.

        In order to be eligible for inclusion in the registry, a sampling
        method's init method must work with just a single argument, a
        Stochastic object.

        If you want to exclude a particular step method from
        consideration for handling a variable, do this:

        Competence functions MUST be called 'competence' and be decorated by the
        '@staticmethod' decorator. Example:

            @staticmethod
            def competence(s):
                if isinstance(s, MyStochasticSubclass):
                    return 2
                else:
                    return 0

        :SeeAlso: pick_best_methods, assign_method
        """
        return 0

    def tune(self, *args, **kwargs):
        return False

    def _get_loglike(self):
        # Fetch log-probability (as sum of childrens' log probability)
        sum = logp_of_set(self.children)
        if self.verbose>1:
            print '\t' + self._id + ' Current log-likelihood ', sum
        return sum

    # Make get property for retrieving log-probability
    loglike = property(fget = _get_loglike, doc="The summed log-probability of all stochastic variables that depend on \n self.stochastics, with self.stochastics removed.")

    def _get_logp_plus_loglike(self):
        sum = logp_of_set(self.markov_blanket)
        if self.verbose>1:
            print '\t' + self._id + ' Current log-likelihood plus current log-probability', sum
        return sum

    # Make get property for retrieving log-probability
    logp_plus_loglike = property(fget = _get_logp_plus_loglike, doc="The summed log-probability of all stochastic variables that depend on \n self.stochastics, and self.stochastics.")

    def current_state(self):
        """Return a dictionary with the current value of the variables defining
        the state of the step method."""
        state = {}
        for s in self._state:
            state[s] = getattr(self, s)
        return state

    @prop
    def ratio():
        """Acceptance ratio"""
        def fget(self):
            return self.accepted/(self.accepted + self.rejected)
        return locals()

class NoStepper(StepMethod):
    """
    Step and tune methods do nothing.

    Useful for holding stochastics constant without setting observed=True.
    """
    def step(self):
        pass
    def tune(self, *args, **kwargs):
        pass

# The default StepMethod, which Model uses to handle singleton stochastics.
class Metropolis(StepMethod):
    """
    The default StepMethod, which Model uses to handle singleton, continuous variables.

    Applies the one-at-a-time Metropolis-Hastings algorithm to the Stochastic over which self has jurisdiction.

    To instantiate a Metropolis called M with jurisdiction over a Stochastic P:

      >>> M = Metropolis(P, scale=1, proposal_sd=None, dist=None)

    :Arguments:
    - s : Stochastic
            The variable over which self has jurisdiction.

    - scale (optional) : number
            The proposal jump width is set to scale * variable.value.

    - proposal_sd (optional) : number or vector
            The proposal jump width is set to proposal_sd.

    - proposal_distribution (optional) : string
            The proposal distribution. May be 'Normal', 'RoundedNormal', 'Bernoulli',
            'Prior' or None. If None is provided, a proposal distribution is chosen
            by examining P.value's type.

    - verbose (optional) : None or integer
            Level of output verbosity: 0=none, 1=low, 2=medium, 3=high. Setting to none allows verbosity to be turned on by sampler.

    :SeeAlso: StepMethod, Sampler.
    """

    def __init__(self, stochastic, scale=1., proposal_sd=None, proposal_distribution=None, verbose=None, tally=True):
        # Metropolis class initialization

        # Initialize superclass
        StepMethod.__init__(self, [stochastic], tally=tally)

        # Initialize hidden attributes
        self.adaptive_scale_factor = 1.
        self.accepted = 0.
        self.rejected = 0.
        self._state = ['rejected', 'accepted', 'adaptive_scale_factor', 'proposal_sd', 'proposal_distribution']
        self._tuning_info = ['adaptive_scale_factor']

        # Set public attributes
        self.stochastic = stochastic
        if verbose is not None:
            self.verbose = verbose
        else:
            self.verbose = stochastic.verbose

        # Avoid zeros when setting proposal variance
        if proposal_sd is not None:
            self.proposal_sd = proposal_sd
        else:
            if all(self.stochastic.value != 0.):
                self.proposal_sd = ones(shape(self.stochastic.value)) * abs(self.stochastic.value) * scale
            else:
                self.proposal_sd = ones(shape(self.stochastic.value)) * scale

        # Initialize proposal deviate with array of zeros
        self.proposal_deviate = zeros(shape(self.stochastic.value), dtype=float)

        # Determine size of stochastic
        if isinstance(self.stochastic.value, ndarray):
            self._len = len(self.stochastic.value.ravel())
        else:
            self._len = 1

        # If no dist argument is provided, assign a proposal distribution automatically.
        if not proposal_distribution:

            # Pick Gaussian by default
            self.proposal_distribution = "Normal"

        else:
            self.proposal_distribution = proposal_distribution

    @staticmethod
    def competence(s):
        """
        The competence function for Metropolis
        """
        if s.dtype is None:
            return .5

        if not s.dtype in float_dtypes:
            # If the stochastic's binary or discrete, I can't do it.
            return 0
        else:
            return 1

    def hastings_factor(self):
        """
        If this is a Metropolis-Hastings method (proposal is not symmetric random walk),
        this method should return log(back_proposal) - log(forward_proposal).
        """
        return 0.

    def step(self):
        """
        The default step method applies if the variable is floating-point
        valued, and is not being proposed from its prior.
        """

        # Probability and likelihood for s's current value:

        if self.verbose>1:
            print
            print self._id + ' getting initial logp.'

        if self.proposal_distribution == "Prior":
            logp = self.loglike
        else:
            logp = self.logp_plus_loglike

        if self.verbose>1:
            print self._id + ' proposing.'

        # Sample a candidate value
        self.propose()

        # Probability and likelihood for s's proposed value:
        try:
            if self.proposal_distribution == "Prior":
                logp_p = self.loglike
                # Check for weirdness before accepting jump
                self.stochastic.logp
            else:
                logp_p = self.logp_plus_loglike

        except ZeroProbability:

            # Reject proposal
            if self.verbose>1:
                print self._id + ' rejecting due to ZeroProbability.'
            self.reject()

            # Increment rejected count
            self.rejected += 1

            if self.verbose>1:
                print self._id + ' returning.'
            return

        if self.verbose>1:
            print 'logp_p - logp: ', logp_p - logp

        HF = self.hastings_factor()

        # Evaluate acceptance ratio
        if log(random()) > logp_p - logp + HF:

            # Revert s if fail
            self.reject()

            # Increment rejected count
            self.rejected += 1
            if self.verbose > 1:
                print self._id + ' rejecting'
        else:
            # Increment accepted count
            self.accepted += 1
            if self.verbose > 1:
                print self._id + ' accepting'

        if self.verbose > 1:
            print self._id + ' returning.'

    def tune(self, *args, **kwargs):
        if self.proposal_distribution == "Prior":
            return False
        else:
            return StepMethod.tune(self, *args, **kwargs)

    def reject(self):
        # Sets current s value to the last accepted value
        # self.stochastic.value = self.stochastic.last_value
        self.stochastic.revert()

    def propose(self):
        """
        This method is called by step() to generate proposed values
        if self.proposal_distribution is "Normal" (i.e. no proposal specified).
        """
        if self.proposal_distribution == "Normal":
            self.stochastic.value = rnormal(self.stochastic.value, self.adaptive_scale_factor * self.proposal_sd)
        elif self.proposal_distribution == "Prior":
            self.stochastic.random()

    def tune(self, divergence_threshold=1e10, verbose=0):
        """
        Tunes the scaling parameter for the proposal distribution
        according to the acceptance rate of the last k proposals:

        Rate    Variance adaptation
        ----    -------------------
        <0.001        x 0.1
        <0.05         x 0.5
        <0.2          x 0.9
        >0.5          x 1.1
        >0.75         x 2
        >0.95         x 10

        This method is called exclusively during the burn-in period of the
        sampling algorithm.

        May be overridden in subclasses.
        """

        if self.verbose is not None:
            verbose = self.verbose

        if self.verbose is not None:
            verbose = self.verbose

        # Verbose feedback
        if verbose > 0:
            print '\t%s tuning:' % self._id

        # Flag for tuning state
        tuning = True

        # Calculate recent acceptance rate
        if not (self.accepted + self.rejected): return tuning
        acc_rate = self.accepted / (self.accepted + self.rejected)

        # Switch statement
        if acc_rate<0.001:
            # reduce by 90 percent
            self.adaptive_scale_factor *= 0.1
        elif acc_rate<0.05:
            # reduce by 50 percent
            self.adaptive_scale_factor *= 0.5
        elif acc_rate<0.2:
            # reduce by ten percent
            self.adaptive_scale_factor *= 0.9
        elif acc_rate>0.95:
            # increase by factor of ten
            self.adaptive_scale_factor *= 10.0
        elif acc_rate>0.75:
            # increase by double
            self.adaptive_scale_factor *= 2.0
        elif acc_rate>0.5:
            # increase by ten percent
            self.adaptive_scale_factor *= 1.1
        else:
            tuning = False

        # Re-initialize rejection count
        self.rejected = 0.
        self.accepted = 0.

        # More verbose feedback, if requested
        if verbose > 0:
            if hasattr(self, 'stochastic'):
                print '\t\tvalue:', self.stochastic.value
            print '\t\tacceptance rate:', acc_rate
            print '\t\tadaptive scale factor:', self.adaptive_scale_factor
            print

        return tuning

class PDMatrixMetropolis(Metropolis):
    """Metropolis sampler with proposals customised for symmetric positive definite matrices"""
    def __init__(self, stochastic, scale=1., proposal_sd=None, verbose=None, tally=True):
        Metropolis.__init__(self, stochastic, scale=scale, proposal_sd=proposal_sd, proposal_distribution="Normal", verbose=verbose, tally=tally)

    @staticmethod
    def competence(s):
        """
        The competence function for MatrixMetropolis
        """
        # MatrixMetropolis handles the Wishart family, which are valued as
        # _symmetric_ matrices.
        if any([isinstance(s,cls) for cls in [distributions.Wishart,distributions.InverseWishart,distributions.WishartCov]]):
            return 2
        else:
            return 0

    def propose(self):
        """
        Proposals for positive definite matrix using random walk deviations on the Cholesky
        factor of the current value.
        """

        # Locally store size of matrix
        dims = self.stochastic.value.shape

        # Add normal deviate to value and symmetrize
        dev =  rnormal(0, self.adaptive_scale_factor * self.proposal_sd, size=dims)
        symmetrize(dev)

        # Replace
        self.stochastic.value = dev + self.stochastic.value

class Gibbs(Metropolis):
    """
    Base class for the Gibbs step methods
    """
    def __init__(self, stochastic, verbose=None):
        Metropolis.__init__(self, stochastic, verbose=verbose, tally=False)

    # Override Metropolis's competence.
    competence = classmethod(StepMethod.competence)

    def step(self):
        if not self.conjugate:
            logp = self.stochastic.logp

        self.propose()

        if not self.conjugate:

            try:
                logp_p = self.stochastic.logp
            except ZeroProbability:
                self.reject()

            if log(np.random.random()) > logp_p - logp:
                self.reject()

    def tune(self, *args, **kwargs):
        return False

    def propose(self):
        raise NotImplementedError, 'The Gibbs class has to be subclassed, it is not usable directly.'

class DrawFromPrior(StepMethod):
    """
    Handles dataless submodels.
    """
    def __init__(self, variables, generations, verbose=None):
        StepMethod.__init__(self, variables, verbose, tally=False)
        self.generations = generations

        # Some variables (eg GP) may not have logp attributes, so don't try to
        # evaluate their logps.
        self.variables_with_logp = set([])
        for s in self.markov_blanket:
            try:
                s.logp
                self.variables_with_logp.add(s)
            except:
                pass

    def get_logp_plus_loglike(self):
        return logp_of_set(self.variables_with_logp)
    logp_plus_loglike = property(get_logp_plus_loglike)

    def step(self):
        jumped = []
        try:
            for generation in self.generations:
                for s in generation:
                    s.rand()
                    jumped.append(s)
            self.logp_plus_loglike
        except ZeroProbability:
            if self.verbose > 0:
                forbidden = []
                for generation in self.generations:
                    for s in self.stochastics:
                        try:
                            s.logp
                        except ZeroProbability:
                            forbidden.append(s.__name__)
                print 'DrawFromPrior jumped stochastics %s to value forbidden by objects %s, rejecting.'%(', '.join(s.__name__ for s in jumped),', '.join(forbidden))
            warnings.warn('DrawFromPrior jumped to forbidden value')
            for s in jumped:
                s.revert()

    @classmethod
    def competence(s):
        # Dataless gets assigned specially before other step methods.
        return 0

class NoStepper(StepMethod):
    """
    Step and tune methods do nothing.

    Useful for holding stochastics constant without setting observed=True.
    """
    def step(self, *args, **kwargs):
        pass
    def tune(self, *args, **kwargs):
        return False

class DiscreteMetropolis(Metropolis):
    """
    Just like Metropolis, but rounds the variable's value.
    Good for discrete stochastics.
    """
    def __init__(self, stochastic, scale=1., proposal_sd=None, proposal_distribution="Poisson", positive=False, verbose=None, tally=True):
        # DiscreteMetropolis class initialization

        # Initialize superclass
        Metropolis.__init__(self, stochastic, scale=scale, proposal_sd=proposal_sd, proposal_distribution=proposal_distribution, verbose=verbose, tally=tally)

        # Flag for positive-only values
        self._positive = positive

    @staticmethod
    def competence(stochastic):
        """
        The competence function for DiscreteMetropolis.
        """
        if stochastic.dtype in integer_dtypes:
            return 1
        else:
            return 0

    def propose(self):
        # Propose new values using normal distribution

        if self.proposal_distribution == "Normal":

            # New normal deviate, centred on current value
            new_val = rnormal(self.stochastic.value, self.adaptive_scale_factor * self.proposal_sd)

            # Round before setting proposed value
            self.stochastic.value = round_array(new_val)

        elif self.proposal_distribution == "Poisson":

            k = shape(self.stochastic.value)
            # Add or subtract (equal probability) Poisson sample
            new_val = self.stochastic.value + rpoisson(self.adaptive_scale_factor * self.proposal_sd) * (-ones(k))**(random(k)>0.5)

            if self._positive:
                # Enforce positive values
                self.stochastic.value = abs(new_val)
            else:
                self.stochastic.value = new_val

        elif self.proposal_distribution == "Prior":
            self.stochastic.random()

class BinaryMetropolis(Metropolis):
    """
    Like Metropolis, but with a modified step() method.
    Good for binary variables.

    """

    def __init__(self, stochastic, p_jump=.1, proposal_distribution=None, verbose=None, tally=True):
        # BinaryMetropolis class initialization

        # Initialize superclass
        Metropolis.__init__(self, stochastic, proposal_distribution=proposal_distribution, verbose=verbose, tally=tally)

        self._state.remove('proposal_sd')

        # adaptive_scale_factor controls the jump probability
        self.adaptive_scale_factor = log(1.-p_jump) / log(.5)

    @staticmethod
    def competence(stochastic):
        """
        The competence function for Binary One-At-A-Time Metropolis
        """
        if stochastic.dtype in bool_dtypes:
            return 1
        else:
            return 0

    def step(self):
        if not isscalar(self.stochastic.value):
            Metropolis.step(self)
        else:

            # See what log-probability of True is.
            self.stochastic.value = True

            try:
                logp_true = self.logp_plus_loglike
            except ZeroProbability:
                self.stochastic.value = False
                return

            # See what log-probability of False is.
            self.stochastic.value = False

            try:
                logp_false = self.logp_plus_loglike
            except ZeroProbability:
                self.stochastic.value = True
                return

            # Test
            p_true = exp(logp_true)
            p_false = exp(logp_false)

            if self.verbose>1:
                print """%s step information:
    - logp_true: %f
    - logp_false: %f
    - p_true: %f
    - p_false: %f
                """ % (self._id, logp_true, logp_false, p_true, p_false)

            # Stochastically set value according to relative
            # probabilities of True and False
            if random() > p_false / (p_true + p_false):
                if self.verbose > 1:
                    print "%s setting %s's value to True." % (self._id, self.stochastic)
                self.stochastic.value = True
            elif self.verbose > 1:
                print "%s setting %s's value to False." % (self._id, self.stochastic)

    def propose(self):
        # Propose new values

        if self.proposal_distribution == 'Prior':
            self.stochastic.random()
        else:
            # Convert adaptive_scale_factor to a jump probability
            p_jump = 1.-.5**self.adaptive_scale_factor

            rand_array = random(size=shape(self.stochastic.value))
            new_value = copy(self.stochastic.value)
            switch_locs = where(rand_array<p_jump)
            new_value[switch_locs] = True - new_value[switch_locs]
            # print switch_locs, rand_array, new_value, self.stochastic.value
            self.stochastic.value = new_value

class AdaptiveMetropolis(StepMethod):
    """
    The AdaptativeMetropolis (AM) sampling algorithm works like a regular
    Metropolis, with the exception that stochastic parameters are block-updated
    using a multivariate jump distribution whose covariance is tuned during
    sampling. Although the chain is non-Markovian, i.e. the proposal
    distribution is asymmetric, it has correct ergodic properties. See
    (Haario et al., 2001) for details.

    :Parameters:
      - stochastic : PyMC objects
          Stochastic objects to be handled by the AM algorith,

      - cov : array
          Initial guess for the covariance matrix C. If it is None, the
          covariance will be estimated using the scales dictionary if provided,
          the existing trace if available, or the current stochastics value.
          It is suggested to provide a sensible guess for the covariance, and
          not rely on the automatic assignment from stochastics value.

      - delay : int
          Number of steps before the empirical covariance is computed. If greedy
          is True, the algorithm waits for delay *accepted* steps before computing
          the covariance.

      - interval : int
          Interval between covariance updates. Higher dimensional spaces require
          more samples to obtain reliable estimates for the covariance updates.

      - greedy : bool
          If True, only the accepted jumps are tallied in the internal trace
          until delay is reached. This is useful to make sure that the empirical
          covariance has a sensible structure.

      - shrink_if_necessary : bool
          If True, the acceptance rate is checked when the step method tunes. If
          the acceptance rate is small, the proposal covariance is shrunk according
          to the following rule:

          if acc_rate < .001:
              self.C *= .01
          elif acc_rate < .01:
              self.C *= .25

      - scales : dict
          Dictionary containing the scale for each stochastic keyed by name.
          If cov is None, those scales are used to define an initial covariance
          matrix. If neither cov nor scale is given, the initial covariance is
          guessed from the trace (it if exists) or the objects value, alt

      - verbose : int
          Controls the verbosity level.

    :Notes:
    Use the methods: `cov_from_scales`, `cov_from_trace` and `cov_from_values` for
    more control on the creation of an initial covariance matrix. A lot of problems
    can be avoided with a good initial covariance and long enough intervals between
    covariance updates. That is, do not compensate for a bad covariance guess by
    reducing the interval between updates thinking the covariance matrix will
    converge more rapidly.

    :Reference:
      Haario, H., E. Saksman and J. Tamminen, An adaptive Metropolis algorithm,
          Bernouilli, vol. 7 (2), pp. 223-242, 2001.
    """
    def __init__(self, stochastic, cov=None, delay=1000, interval=200, greedy=True, shrink_if_necessary=False, scales=None, verbose=None, tally=False):

        # Verbosity flag
        self.verbose = verbose

        self.accepted = 0
        self.rejected = 0

        if not np.iterable(stochastic) or isinstance(stochastic, Variable):
            stochastic = [stochastic]

        # Initialize superclass
        StepMethod.__init__(self, stochastic, verbose, tally)

        self._id = 'AdaptiveMetropolis_'+'_'.join([p.__name__ for p in self.stochastics])
        # State variables used to restore the state in a latter session.
        self._state += ['accepted', 'rejected', '_trace_count', '_current_iter', 'C', 'proposal_sd',
        '_proposal_deviate', '_trace', 'shrink_if_necessary']
        self._tuning_info = ['C']

        self.proposal_sd = None
        self.shrink_if_necessary=shrink_if_necessary

        # Number of successful steps before the empirical covariance is computed
        self.delay = delay
        # Interval between covariance updates
        self.interval = interval
        # Flag for tallying only accepted jumps until delay reached
        self.greedy = greedy

        # Initialization methods
        self.check_type()
        self.dimension()

        # Set the initial covariance using cov, or the following fallback mechanisms:
        # 1. If scales is provided, use it.
        # 2. If a trace is present, compute the covariance matrix empirically from it.
        # 3. Use the stochastics value as a guess of the variance.
        if cov is not None:
            self.C = cov
        elif scales:
            self.C = self.cov_from_scales(scales)
        else:
            try:
                self.C = self.cov_from_trace()
            except AttributeError:
                self.C = self.cov_from_value(100.)

        self.updateproposal_sd()

        # Keep track of the internal trace length
        # It may be different from the iteration count since greedy
        # sampling can be done during warm-up period.
        self._trace_count = 0
        self._current_iter = 0

        self._proposal_deviate = np.zeros(self.dim)
        self.chain_mean = np.asmatrix(np.zeros(self.dim))
        self._trace = []

        if self.verbose >= 1:
            print "Initialization..."
            print 'Dimension: ', self.dim
            print "C_0: ", self.C
            print "Sigma: ", self.proposal_sd

    @staticmethod
    def competence(stochastic):
        """
        The competence function for AdaptiveMetropolis.
        The AM algorithm is well suited to deal with multivariate
        parameters.
        """
        if not stochastic.dtype in float_dtypes and not stochastic.dtype in integer_dtypes:
            return 0
            # Algorithm is not well-suited to sparse datasets. Dont use if less than
            # 25 percent of values are nonzero
        if np.alen(stochastic.value) == 1:
            return 0
        elif np.alen(stochastic.value) < 5:
            return 2
        elif (len(stochastic.value.nonzero()[0]) > 0.25*len(stochastic.value)):
            return 2
        else:
            return 0

    def cov_from_value(self, scaling):
        """Return a covariance matrix for the jump distribution using
        the actual value of the stochastic as a guess of their variance,
        divided by the `scaling` argument.

        Note that this is likely to return a poor guess.
        """
        rv = []
        for s in self.stochastics:
            rv.extend(np.ravel(s.value).copy())

        # Remove 0 values since this would lead to quite small jumps...
        arv = np.array(rv)
        arv[arv==0] = 1.

        # Create a diagonal covariance matrix using the scaling factor.
        return np.eye(self.dim)*np.abs(arv)/scaling

    def cov_from_scales(self, scales):
        """Return a covariance matrix built from a dictionary of scales.

        `scales` is a dictionary keyed by stochastic instances, and the
        values refer are the variance of the jump distribution for each
        stochastic. If a stochastic is a sequence, the variance must
        have the same length.
        """

        # Get array of scales
        ord_sc = []
        for stochastic in self.stochastics:
            ord_sc.append(np.ravel(scales[stochastic]))
        ord_sc = np.concatenate(ord_sc)

        if np.squeeze(ord_sc).shape[0] != self.dim:
            raise "Improper initial scales, dimension don't match", \
                (np.squeeze(ord_sc), self.dim)

        # Scale identity matrix
        return np.eye(self.dim)*ord_sc

    def cov_from_trace(self, trace=slice(None)):
        """Define the jump distribution covariance matrix from the object's
        stored trace.

        :Parameters:
        - `trace` : slice or int
          A slice for the stochastic object's trace in the last chain, or a
          an integer indicating the how many of the last samples will be used.

        """
        n = []
        for s in self.stochastics:
            n.append(s.trace.length())
        n = set(n)
        if len(n) > 1:
            raise ValueError, 'Traces do not have the same length.'
        elif n == 0:
            raise AttributeError, 'Stochastic has no trace to compute covariance.'
        else:
            n = n.pop()

        if type(trace) is not slice:
            trace = slice(trace, n)

        a = self.trace2array(trace)
        return np.cov(a, rowvar=0)

    def check_type(self):
        """Make sure each stochastic has a correct type, and identify discrete stochastics."""
        self.isdiscrete = {}
        for stochastic in self.stochastics:
            if stochastic.dtype in integer_dtypes:
                self.isdiscrete[stochastic] = True
            elif stochastic.dtype in bool_dtypes:
                raise 'Binary stochastics not supported by AdaptativeMetropolis.'
            else:
                self.isdiscrete[stochastic] = False

    def dimension(self):
        """Compute the dimension of the sampling space and identify the slices
        belonging to each stochastic.
        """
        self.dim = 0
        self._slices = {}
        for stochastic in self.stochastics:
            if isinstance(stochastic.value, np.matrix):
                p_len = len(stochastic.value.A.ravel())
            elif isinstance(stochastic.value, np.ndarray):
                p_len = len(stochastic.value.ravel())
            else:
                p_len = 1
            self._slices[stochastic] = slice(self.dim, self.dim + p_len)
            self.dim += p_len

    def update_cov(self):
        """Recursively compute the covariance matrix for the multivariate normal
        proposal distribution.

        This method is called every self.interval once self.delay iterations
        have been performed.
        """

        scaling = (2.4)**2/self.dim # Gelman et al. 1996.
        epsilon = 1.0e-5
        chain = np.asarray(self._trace)

        # Recursively compute the chain mean
        self.C, self.chain_mean = self.recursive_cov(self.C, self._trace_count,
            self.chain_mean, chain, scaling=scaling, epsilon=epsilon)

        # Shrink covariance if acceptance rate is too small
        acc_rate = self.accepted / (self.accepted + self.rejected)
        if self.shrink_if_necessary:
            if acc_rate < .001:
                self.C *= .01
            elif acc_rate < .01:
                self.C *= .25
            if self.verbose > 0:
                if acc_rate < .01:
                    print '\tAcceptance rate was',acc_rate,'shrinking covariance'
        self.accepted = 0.
        self.rejected = 0.

        if self.verbose > 0:
            print "\tUpdating covariance ...\n", self.C
            print "\tUpdating mean ... ", self.chain_mean

        # Update state
        adjustmentwarning = '\n'+\
        'Covariance was not positive definite and proposal_sd cannot be computed by \n'+ \
        'Cholesky decomposition. The next jumps will be based on the last \n' + \
        'valid covariance matrix. This situation may have arisen because no \n' + \
        'jumps were accepted during the last `interval`. One solution is to \n' + \
        'increase the interval, or specify an initial covariance matrix with \n' + \
        'a smaller variance. For this simulation, each time a similar error \n' + \
        'occurs, proposal_sd will be reduced by a factor .9 to reduce the \n' + \
        'jumps and increase the likelihood of accepted jumps.'

        try:
            self.updateproposal_sd()
        except np.linalg.LinAlgError:
            warnings.warn(adjustmentwarning)
            self.covariance_adjustment(.9)

        self._trace_count += len(self._trace)
        self._trace = []

    def covariance_adjustment(self, f=.9):
        """Multiply self.proposal_sd by a factor f. This is useful when the current proposal_sd is too large and all jumps are rejected.
        """
        self.proposal_sd *= f

    def updateproposal_sd(self):
        """Compute the Cholesky decomposition of self.C."""
        self.proposal_sd = np.linalg.cholesky(self.C)

    def recursive_cov(self, cov, length, mean, chain, scaling=1, epsilon=0):
        r"""Compute the covariance recursively.

        Return the new covariance and the new mean.

        .. math::
            C_k & = \frac{1}{k-1} (\sum_{i=1}^k x_i x_i^T - k\bar{x_k}\bar{x_k}^T)
            C_n & = \frac{1}{n-1} (\sum_{i=1}^k x_i x_i^T + \sum_{i=k+1}^n x_i x_i^T - n\bar{x_n}\bar{x_n}^T)
                & = \frac{1}{n-1} ((k-1)C_k + k\bar{x_k}\bar{x_k}^T + \sum_{i=k+1}^n x_i x_i^T - n\bar{x_n}\bar{x_n}^T)

        :Parameters:
            -  cov : matrix
                Previous covariance matrix.
            -  length : int
                Length of chain used to compute the previous covariance.
            -  mean : array
                Previous mean.
            -  chain : array
                Sample used to update covariance.
            -  scaling : float
                Scaling parameter
            -  epsilon : float
                Set to a small value to avoid singular matrices.
        """
        n = length + len(chain)
        k = length
        new_mean = self.recursive_mean(mean, length, chain)

        t0 = k * np.outer(mean, mean)
        t1 = np.dot(chain.T, chain)
        t2 = n*np.outer(new_mean, new_mean)
        t3 = epsilon * np.eye(cov.shape[0])

        new_cov =  (k-1)/(n-1.)*cov + scaling/(n-1.) * (t0 + t1 - t2 + t3)
        return new_cov, new_mean

    def recursive_mean(self, mean, length, chain):
        r"""Compute the chain mean recursively.

        Instead of computing the mean :math:`\bar{x_n}` of the entire chain,
        use the last computed mean :math:`bar{x_j}` and the tail of the chain
        to recursively estimate the mean.

        .. math::
            \bar{x_n} & = \frac{1}{n} \sum_{i=1}^n x_i
                      & = \frac{1}{n} (\sum_{i=1}^j x_i + \sum_{i=j+1}^n x_i)
                      & = \frac{j\bar{x_j}}{n} + \frac{\sum_{i=j+1}^n x_i}{n}

        :Parameters:
            -  mean : array
                Previous mean.
            -  length : int
                Length of chain used to compute the previous mean.
            -  chain : array
                Sample used to update mean.
        """
        n = length + len(chain)
        return length * mean / n + chain.sum(0)/n

    def propose(self):
        """
        This method proposes values for stochastics based on the empirical
        covariance of the values sampled so far.

        The proposal jumps are drawn from a multivariate normal distribution.
        """

        arrayjump = np.dot(self.proposal_sd, np.random.normal(size=self.proposal_sd.shape[0]))
        if self.verbose > 2:
            print 'Jump :', arrayjump

        # Update each stochastic individually.
        for stochastic in self.stochastics:
            jump = arrayjump[self._slices[stochastic]]
            if np.iterable(stochastic.value):
                jump = np.reshape(arrayjump[self._slices[stochastic]],np.shape(stochastic.value))
            if self.isdiscrete[stochastic]:
                jump = round_array(jump)
            stochastic.value = stochastic.value + jump

    def step(self):
        """
        Perform a Metropolis step.

        Stochastic parameters are block-updated using a multivariate normal
        distribution whose covariance is updated every self.interval once
        self.delay steps have been performed.

        The AM instance keeps a local copy of the stochastic parameter's trace.
        This trace is used to computed the empirical covariance, and is
        completely independent from the Database backend.

        If self.greedy is True and the number of iterations is smaller than
        self.delay, only accepted jumps are stored in the internal
        trace to avoid computing singular covariance matrices.
        """

        # Probability and likelihood for stochastic's current value:
        logp = self.logp_plus_loglike
        if self.verbose > 1:
            print 'Current value: ', self.stoch2array()
            print 'Current likelihood: ', logp

        # Sample a candidate value
        self.propose()

        # Metropolis acception/rejection test
        accept = False
        try:
            # Probability and likelihood for stochastic's proposed value:
            logp_p = self.logp_plus_loglike
            if self.verbose > 2:
                print 'Current value: ', self.stoch2array()
                print 'Current likelihood: ', logp

            if np.log(random()) < logp_p - logp:
                accept = True
                self.accepted += 1
                if self.verbose > 2:
                    print 'Accepted'
            else:
                self.rejected += 1
                if self.verbose > 2:
                    print 'Rejected'
        except ZeroProbability:
            self.rejected += 1
            logp_p = None
            if self.verbose > 2:
                    print 'Rejected with ZeroProbability Error.'

        if (not self._current_iter % self.interval) and self.verbose > 1:
            print "Step ", self._current_iter
            print "\tLogprobability (current, proposed): ", logp, logp_p
            for stochastic in self.stochastics:
                print "\t", stochastic.__name__, stochastic.last_value, stochastic.value
            if accept:
                print "\tAccepted\t*******\n"
            else:
                print "\tRejected\n"
            print "\tAcceptance ratio: ", self.accepted/(self.accepted+self.rejected)

        if self._current_iter == self.delay:
            self.greedy = False

        if not accept:
            self.reject()

        if accept or not self.greedy:
            self.internal_tally()

        if self._current_iter>self.delay and self._current_iter%self.interval==0:
           self.update_cov()

        self._current_iter += 1

    # Please keep reject() factored out- helps RandomRealizations figure out what to do.
    def reject(self):
        for stochastic in self.stochastics:
            # stochastic.value = stochastic.last_value
            stochastic.revert()

    def internal_tally(self):
        """Store the trace of stochastics for the computation of the covariance.
        This trace is completely independent from the backend used by the
        sampler to store the samples."""
        chain = []
        for stochastic in self.stochastics:
            chain.append(np.ravel(stochastic.value))
        self._trace.append(np.concatenate(chain))

    def trace2array(self, sl):
        """Return an array with the trace of all stochastics, sliced by sl."""
        chain = []
        for stochastic in self.stochastics:
            tr = stochastic.trace.gettrace(slicing=sl)
            chain.append(tr)
        return np.hstack(chain)

    def stoch2array(self):
        """Return the stochastic objects as an array."""
        a = np.empty(self.dim)
        for stochastic in self.stochastics:
            a[self._slices[stochastic]] = stochastic.value
        return a

    def tune(self, verbose=0):
        """Tuning is done during the entire run, independently from the Sampler
        tuning specifications. """
        return False

class IIDSStepper(StepMethod):
    """
    See written documentation.
    """
    pass


"""
@author:  Yuhao Cheng
@contact: yuhao.cheng[at]outlook.com
"""
#!!!!! ignore the warning messages
import warnings
warnings.filterwarnings('ignore')
import os
import pickle
import math
import torch
import time
import numpy as np
from PIL import Image
from collections import OrderedDict
import torchvision.transforms as T
import torchvision.transforms.functional as tf
from torch.utils.data import DataLoader

import logging
logger = logging.getLogger(__name__)

from pyanomaly.core.utils import AverageMeter, flow_batch_estimate, tensorboard_vis_images, vis_optical_flow, make_info_message, ParamSet
from pyanomaly.datatools.evaluate.utils import psnr_error

from pyanomaly.datatools.evaluate.utils import (
    simple_diff,
    find_max_patch,
    amc_score,
    calc_w
    )

from ..abstract.base_engine import BaseTrainer, BaseInference, BaseService

from ..engine_registry import ENGINE_REGISTRY

__all__ = ['MATrainer', 'AMCInference']

@ENGINE_REGISTRY.register()
class MATrainer(BaseTrainer):
    """
    G
    D_frame
    D_pattern
    AE_act
    AE_obj
    PatternNet
    """
    NAME = ["MA.TRAIN"]
    def custom_setup(self):
        # create loss meters
        self.loss_meter_G = AverageMeter(name='Loss_G')
        self.loss_meter_D = AverageMeter(name='Loss_D')

        self.optical = ParamSet(name='optical', size=self.config.DATASET.optical_size, output_format=self.config.DATASET.optical_format)
        # import ipdb; ipdb.set_trace()

    def train(self,current_step):
        # Pytorch [N, C, D, H, W]
        # initialize
        start = time.time()
        self.set_requires_grad(self.F, False)
        self.set_requires_grad(self.D, True)
        self.set_requires_grad(self.G, True)
        self.G.train()
        self.D.train()
        self.F.eval()
        writer = self.kwargs['writer_dict']['writer']
        global_steps = self.kwargs['writer_dict']['global_steps_{}'.format(self.kwargs['model_type'])]

        # get the data
        data, anno, meta = next(self._train_loader_iter)
        self.data_time.update(time.time() - start)

        # base on the D to get each frame
        # in this method, D = 2 and not change
        input_data = data[:, :, 0, :, :].cuda() # input(1-st) frame
        target = data[:, :, 1,:, :].cuda() # target(2-nd) frame

        # True Process =================Start===================
        #---------update optim_G ---------
        self.set_requires_grad(self.D, False)
        output_flow_G,  output_frame_G = self.G(input_data)
        gt_flow_esti_tensor = torch.cat([input_data, target], 1)
        flow_gt_vis, flow_gt  = flow_batch_estimate(self.F, gt_flow_esti_tensor, self.normalize.param['train'],
                                                    optical_size=self.config.DATASET.optical_size, output_format=self.config.DATASET.optical_format)
        fake_g = self.D(torch.cat([target, output_flow_G], dim=1))

        loss_g_adv = self.GANLoss(fake_g, True)
        loss_op = self.OpticalflowSqrtLoss(output_flow_G, flow_gt)
        loss_int = self.IntentsityLoss(output_frame_G, target)
        loss_gd = self.GradientLoss(output_frame_G, target)
        loss_g_all = self.loss_lamada['IntentsityLoss'] * loss_int + self.loss_lamada['GradientLoss'] * loss_gd + self.loss_lamada['OpticalflowSqrtLoss'] * loss_op + self.loss_lamada['GANLoss'] * loss_g_adv

        self.optimizer_G.zero_grad()
        loss_g_all.backward()
        self.optimizer_G.step()
        self.loss_meter_G.update(loss_g_all.detach())

        if self.config.TRAIN.adversarial.scheduler.use:
            self.optimizer_G_scheduler.step()

        #---------update optim_D ---------------
        self.set_requires_grad(self.D, True)
        self.optimizer_D.zero_grad()
        # import ipdb; ipdb.set_trace()
        real_d = self.D(torch.cat([target, flow_gt],dim=1))
        fake_d = self.D(torch.cat([target, output_flow_G.detach()], dim=1))
        loss_d_1 = self.GANLoss(real_d, True)
        loss_d_2 = self.GANLoss(fake_d, False)
        loss_d = (loss_d_1  + loss_d_2) * 0.5
        loss_d.backward()
        self.optimizer_D.step()
        if self.config.TRAIN.adversarial.scheduler.use:
            self.optimizer_D_scheduler.step()
        self.loss_meter_D.update(loss_d.detach())
        # ======================End==================

        self.batch_time.update(time.time() - start)

        if (current_step % self.steps.param['log'] == 0):
            msg = make_info_message(current_step, self.steps.param['max'], self.kwargs['model_type'], self.batch_time,
                                    self.config.TRAIN.batch_size, self.data_time, [self.loss_meter_G, self.loss_meter_D])
            logger.info(msg)

        writer.add_scalar('Train_loss_G', self.loss_meter_G.val, global_steps)
        writer.add_scalar('Train_loss_D', self.loss_meter_D.val, global_steps)

        if (current_step % self.steps.param['vis'] == 0):
            temp = vis_optical_flow(output_flow_G.detach(), output_format=self.config.DATASET.optical_format, output_size=(output_flow_G.shape[-2], output_flow_G.shape[-1]),
                                    normalize=self.normalize.param['train'])
            vis_objects = OrderedDict({
                'train_target_flow': flow_gt_vis.detach(),
                'train_output_flow_G': temp,
                'train_target_frame': target.detach(),
                'train_output_frame_G': output_frame_G.detach(),
            })
            tensorboard_vis_images(vis_objects, writer, global_steps, self.normalize.param['train'])
        global_steps += 1

        # reset start
        start = time.time()

        # self.saved_model = {'G':self.G, 'D':self.D}
        self.saved_model['G'] = self.G
        self.saved_model['D'] = self.D
        # self.saved_optimizer = {'optim_G': self.optimizer_G, 'optim_D': self.optimizer_D}
        self.saved_optimizer['optimizer_G'] = self.optimizer_G
        self.saved_optimizer['optimizer_D'] = self.optimizer_D
        # self.saved_loss = {'loss_G':self.loss_meter_G.val, 'loss_D':self.loss_meter_D.val}
        self.saved_loss['loss_G'] = self.loss_meter_G.val
        self.saved_loss['loss_D'] = self.loss_meter_D.val
        self.kwargs['writer_dict']['global_steps_{}'.format(self.kwargs['model_type'])] = global_steps

@ENGINE_REGISTRY.register()
class AMCInference(BaseInference):
    NAME = ["AMC.INFERENCE"]

    def inference(self):
        for h in self._hooks:
            h.inference()

@ENGINE_REGISTRY.register()
class AMCService(BaseService):
    def custom_setup(self):
        self.optical_format = self.config.DATASET.optical_format
        self.optical_szie = self.engine.config.DATASET.optical_size
        self.wf = 1.0
        self.wi = 1.0
        self.threshold = 0.0 # the threshold to judge whether the frame is the anomaly

    def get_clip_by_stride(self, video, stride=2):
        """Get the clip list by the stride
        """
        return []

    def execute(self, data):
        output_dict = OrderedDict()
        # data.shape = [N,C,D,H,W], data is a whole vide, D=the length of the video
        clip_list = self.get_clip_by_stride(data) # the length of the length is the length of the video
        scores = np.empty(shape=(len(clip_list), ), dtype=np.float32)

        for index, clip in enumerate(clip_list):
            first_frame = clip[:, :, 0, :, :].cuda()
            second_frame = clip[:, :, 1, :, :].cuda()

            generated_flow, generated_frame = self.G(first_frame)
            gtFlowEstim = torch.cat([first_frame, second_frame], 1)
            _, gtFlow = flow_batch_estimate(self.F, gtFlowEstim, self.normalize.param['val'], output_format=self.optical_format, optical_size=self.optical_size)

            score, _, _ = amc_score(second_frame, generated_frame, gtFlow, generated_flow, self.wf, self.wi)
            score = score.tolist()
            scores[index] = score

        result_mask = scores.gt(self.threshold)
        output_dict['result_dict'] = result_mask

        return output_dict


from typing import *

import numpy as np
import torch
from scipy import linalg as la

from ...settings_ import settings
from ...typing_ import TensorInitArgType
from . import init
from .core import *
from .layers import *
from .linalg import *
from .nn import *

__all__ = [
    'Flow', 'FeatureMappingFlow',
    'InverseFlow', 'SequentialFlow',
    'LooseInvertibleMatrix', 'StrictInvertibleMatrix',
    'InvertibleDense', 'InvertibleConv1d', 'InvertibleConv2d',
    'InvertibleConv3d',
    'Scale', 'SigmoidScale', 'ExpScale', 'LinearScale',
]

# ---- base flow classes ----
class BaseValidateTensorLayer(BaseLayer):

    __constants__ = ('_should_validate_tensor', '_validate_tensor_messgae_prefix')

    _should_validate_tensor: bool
    _validate_tensor_messgae_prefix: str

    def __init__(self):
        super().__init__()
        self._should_validate_tensor = bool(settings.validate_tensors)
        self._validate_tensor_messgae_prefix = self.__class__.__qualname__

    @jit_method
    def _maybe_assert_finite(self,
                             t: Tensor,
                             name: str,
                             inverse: bool = False) -> Tensor:
        if self._should_validate_tensor:
            msg = '{}.{}'.format(self._validate_tensor_messgae_prefix, name)
            if inverse:
                msg += ' [inverse]'
            t = assert_finite(t, msg)
        return t

class Flow(BaseValidateTensorLayer):
    """
    Base class for normalizing flows.

    A normalizing flow transforms a random variable `x` into `y` by an
    (implicitly) invertible mapping :math:`y = f(x)`, whose Jaccobian matrix
    determinant :math:`\\det \\frac{\\partial f(x)}{\\partial x} \\neq 0`, thus
    can derive :math:`\\log p(y)` from given :math:`\\log p(x)`.
    """

    __constants__ = BaseValidateTensorLayer.__constants__ + (
        'x_event_ndims', 'y_event_ndims', 'explicitly_invertible'
    )

    x_event_ndims: int
    """Number of event dimensions in `x`."""

    y_event_ndims: int
    """Number of event dimensions in `y`."""

    explicitly_invertible: bool
    """
    Whether or not this flow is explicitly invertible?

    If a flow is not explicitly invertible, then it only supports to
    transform `x` into `y`, and corresponding :math:`\\log p(x)` into
    :math:`\\log p(y)`.  It cannot compute :math:`\\log p(y)` directly
    without knowing `x`, nor can it transform `x` back into `y`.
    """

    def __init__(self,
                 x_event_ndims: int,
                 y_event_ndims: int,
                 explicitly_invertible: bool):
        super().__init__()

        self.x_event_ndims = int(x_event_ndims)
        self.y_event_ndims = int(y_event_ndims)
        self.explicitly_invertible = bool(explicitly_invertible)

    @jit_method
    def get_x_event_ndims(self) -> int:
        return self.x_event_ndims

    @jit_method
    def get_y_event_ndims(self) -> int:
        return self.y_event_ndims

    @jit_method
    def is_explicitly_invertible(self) -> bool:
        return self.explicitly_invertible

    def invert(self) -> 'Flow':
        """
        Get the inverse flow from this flow.

        Specifying `inverse = True` when calling the inverse flow will be
        interpreted as having `inverse = False` in the original flow, and
        vise versa.

        If the current flow requires to be initialized by calling it
        with `inverse = False`, then the inversed flow will require to be
        initialized by calling it with `inverse = True`, and vise versa.

        Returns:
            The inverse flow.
        """
        return InverseFlow(self)

    def _transform(self,
                   input: Tensor,
                   input_log_det: Optional[Tensor],
                   inverse: bool,
                   compute_log_det: bool
                   ) -> Tuple[Tensor, Optional[Tensor]]:
        raise NotImplementedError()

    def forward(self,
                input: Tensor,
                input_log_det: Optional[Tensor] = None,
                inverse: bool = False,
                compute_log_det: bool = True
                ) -> Tuple[Tensor, Optional[Tensor]]:
        """
        Transform `x` into `y` and compute the log-determinant of `f` at `x`
        (if `inverse` is False); or transform `y` into `x` and compute the
        log-determinant of `f^{-1}` at `y` (if `inverse` is True).

        Args:
            input: `x` (if `inverse` is False) or `y` (if `inverse` is True).
            input_log_det: The log-determinant of the previous layer.
                Will add the log-determinant of this layer to `input_log_det`,
                to obtain the output log-determinant.  If no previous layer,
                will start from zero log-det.
            inverse: See above.
            compute_log_det: Whether or not to compute the log-determinant?

        Returns:
            The transformed tensor, and the summed log-determinant of
            the previous flow layer and this layer.
        """
        if inverse:
            if not self.explicitly_invertible:
                raise RuntimeError('Flow is not explicitly invertible.')
            event_ndims = self.y_event_ndims
        else:
            event_ndims = self.x_event_ndims

        if input.dim() < event_ndims:
            raise ValueError(
                '`input` is required to be at least {}d, but the input shape '
                'is {}.'.format(event_ndims, shape(input))
            )

        input_shape = input.shape
        log_det_shape = input_shape[: len(input_shape) - event_ndims]

        if input_log_det is not None:
            if input_log_det.shape != log_det_shape:
                raise ValueError(
                    'The shape of `input_log_det` is not expected: '
                    'expected to be {}, but got {}.'.
                    format(list(log_det_shape), shape(input_log_det))
                )

        # compute the transformed output and log-det
        output, output_log_det = self._transform(
            input, input_log_det, inverse, compute_log_det)

        if output_log_det is not None:
            if output_log_det.shape != log_det_shape:
                output_log_det = output_log_det + torch.zeros(
                    log_det_shape, dtype=output_log_det.dtype, device=output_log_det.device)

            if output_log_det.shape != log_det_shape:
                raise ValueError(
                    'The shape of `output_log_det` is not expected: '
                    'expected to be {}, but got {}.'.
                    format(list(log_det_shape), shape(output_log_det))
                )

        return output, output_log_det

class FeatureMappingFlow(Flow):
    """Base class for flows mapping input features to output features."""

    __constants__ = Flow.__constants__ + ('axis',)

    axis: int
    """The feature axis (negative index)."""

    def __init__(self,
                 axis: int,
                 event_ndims: int,
                 explicitly_invertible: bool):
        """
        Construct a new :class:`FeatureMappingFlow`.

        Args:
            axis: The feature axis, on which to apply the transformation.
                It must be a negative integer, and included in the
                event dimensions.
            event_ndims: Number of event dimensions in both `x` and `y`.
                `x.ndims - event_ndims == log_det.ndims` and
                `y.ndims - event_ndims == log_det.ndims`.
            explicitly_invertible: Whether or not this flow is explicitly
                invertible?
        """
        # check the arguments
        axis = int(axis)
        event_ndims = int(event_ndims)

        if event_ndims < 1:
            raise ValueError(f'`event_ndims` must be at least 1: '
                             f'got {event_ndims}')

        if axis >= 0 or axis < -event_ndims:
            raise ValueError(
                f'`-event_ndims <= axis < 0` does not hold: '
                f'`axis` is {axis}, while `event_ndims` is {event_ndims}.')

        # construct the layer
        super().__init__(x_event_ndims=event_ndims,
                         y_event_ndims=event_ndims,
                         explicitly_invertible=explicitly_invertible)
        self.axis = axis

    @jit_method
    def get_axis(self) -> int:
        return self.axis

    @jit_method
    def get_event_ndims(self) -> int:
        """Get the number of event dimensions in both `x` and `y`."""
        return self.x_event_ndims

# ---- composite flows ----
class InverseFlow(Flow):
    """A flow that inverts another given flow."""

    original_flow: Module
    """The original flow, which is inverted by this :class:`InverseFlow`."""

    def __init__(self, flow: Module):
        if (not isinstance(flow, Flow) and not is_jit_layer(flow)) or \
                not flow.is_explicitly_invertible():
            raise TypeError(
                f'`flow` must be an explicitly invertible flow: '
                f'got {flow!r}'
            )

        super().__init__(
            x_event_ndims=flow.get_y_event_ndims(),
            y_event_ndims=flow.get_x_event_ndims(),
            explicitly_invertible=flow.is_explicitly_invertible(),
        )
        self.original_flow = flow

    def invert(self) -> Flow:
        return self.original_flow

    def _transform(self,
                   input: Tensor,
                   input_log_det: Optional[Tensor],
                   inverse: bool,
                   compute_log_det: bool) -> Tuple[Tensor, Optional[Tensor]]:
        return self.original_flow(
            input, input_log_det, not inverse, compute_log_det)

class _NotInvertibleFlow(Module):

    def forward(self,
                input: Tensor,
                input_log_det: Optional[Tensor],
                inverse: bool,
                compute_log_det: bool
                ) -> Tuple[Tensor, Optional[Tensor]]:
        raise RuntimeError('Not an explicitly invertible flow.')

class SequentialFlow(Flow):

    _chain: ModuleList
    _inverse_chain: ModuleList

    def custom_compile_children(self):
        flows = [jit_compile(m) for m in self._chain]
        self._chain = torch.jit.script(ModuleList(flows))
        self._inverse_chain = torch.jit.script(ModuleList(flows[::-1]))

    def __init__(self,
                 *flows: Union[Module, Sequence[Module]]):
        from tensorkit.layers import flatten_nested_layers

        # validate the arguments
        flows = flatten_nested_layers(flows)
        if not flows:
            raise ValueError('`flows` must not be empty.')

        for i, flow in enumerate(flows):
            if not isinstance(flow, Flow) and not is_jit_layer(flow):
                raise TypeError(f'`flows[{i}]` is not a flow: got {flow!r}')

        for i, (flow1, flow2) in enumerate(zip(flows[:-1], flows[1:])):
            if flow2.get_x_event_ndims() != flow1.get_y_event_ndims():
                raise ValueError(
                    f'`x_event_ndims` of `flows[{i + 1}]` != '
                    f'`y_event_ndims` of `flows[{i}]`: '
                    f'{flow2.get_x_event_ndims()} vs {flow1.get_y_event_ndims()}.'
                )

        super().__init__(
            x_event_ndims=flows[0].get_x_event_ndims(),
            y_event_ndims=flows[-1].get_y_event_ndims(),
            explicitly_invertible=all(
                flow.is_explicitly_invertible() for flow in flows)
        )
        self._chain = ModuleList(flows)
        self._inverse_chain = ModuleList(flows[::-1])

    def _call_chain(self,
                            output: Tensor,
                            output_log_det: Optional[Tensor],
                            compute_log_det: bool
                            ) -> Tuple[Tensor, Optional[Tensor]]:
        for flow in self._chain:
            output, output_log_det = flow(
                output, output_log_det, False, compute_log_det)
        return output, output_log_det

    def _call_inverse_chain(self,
                            output: Tensor,
                            output_log_det: Optional[Tensor],
                            compute_log_det: bool
                            ) -> Tuple[Tensor, Optional[Tensor]]:
        for flow in self._inverse_chain:
            output, output_log_det = flow(
                output, output_log_det, True, compute_log_det)
        return output, output_log_det

    def _transform(self,
                   input: Tensor,
                   input_log_det: Optional[Tensor],
                   inverse: bool,
                   compute_log_det: bool
                   ) -> Tuple[Tensor, Optional[Tensor]]:
        output, output_log_det = input, input_log_det
        event_ndims = self.y_event_ndims if inverse else self.x_event_ndims

        if rank(output) > event_ndims + 1:
            output, batch_shape = flatten_to_ndims(output, event_ndims + 1)
            if output_log_det is not None:
                output_log_det = reshape(output_log_det, [-1])
        else:
            batch_shape: Optional[List[int]] = None

        if inverse:
            output, output_log_det = self._call_inverse_chain(
                output, output_log_det, compute_log_det)
        else:
            output, output_log_det = self._call_chain(
                output, output_log_det, compute_log_det)

        if batch_shape is not None:
            output = unflatten_from_ndims(output, batch_shape)
            if output_log_det is not None:
                output_log_det = reshape(output_log_det, batch_shape)

        return output, output_log_det

# ---- invertible linear flows ----
class InvertibleMatrix(BaseValidateTensorLayer):

    __constants__ = BaseValidateTensorLayer.__constants__ + (
        'size', 'validate_tensors'
    )

    size: int

    validate_tensors: bool
    """Whether or not to perform time-consuming validations on tensors?"""

    def __init__(self, size: int):
        super().__init__()
        self.size = size

        # TODO: make validate_tensors an argument
        self.validate_tensors = settings.validate_tensors is True

    def __repr__(self):
        return f'{self.__class__.__qualname__}(size={self.size})'

class LooseInvertibleMatrix(InvertibleMatrix):
    """
    A matrix initialized to be an invertible, orthogonal matrix.

    There is no guarantee that the matrix will keep invertible during training.
    But according to the measure theory, the non-invertible n by n real matrices
    are of measure 0.  Thus this class is generally enough for use.
    """

    def __init__(self,
                 seed_matrix: np.ndarray,
                 dtype: str = settings.float_x,
                 device: Optional[str] = None):
        """
        Construct a new :class:`LooseInvertibleMatrix`.

        Args:
            seed_matrix: A matrix that is used as a seed to obtain the
                initial invertible and orthogonal matrix.
            dtype: The dtype of the matrix.
            device: The device where to place new tensors and variables.
        """
        device = device or current_device()
        initial_matrix = la.qr(seed_matrix)[0]

        super().__init__(initial_matrix.shape[0])
        add_parameter(
            self, 'matrix',
            from_numpy(initial_matrix, dtype=dtype, device=device)
        )

    def forward(self,
                inverse: bool,
                compute_log_det: bool
                ) -> Tuple[Tensor, Optional[Tensor]]:
        log_det: Optional[Tensor] = None
        if inverse:
            matrix = self._maybe_assert_finite(
                matrix_inverse(self.matrix), 'matrix', inverse)
            if compute_log_det:
                log_det = self._maybe_assert_finite(
                    -slogdet(self.matrix)[1], 'log_det', inverse)
        else:
            matrix = self._maybe_assert_finite(self.matrix, 'matrix', inverse)
            if compute_log_det:
                log_det = self._maybe_assert_finite(
                    slogdet(self.matrix)[1], 'log_det', inverse)

        return matrix, log_det

class StrictInvertibleMatrix(InvertibleMatrix):
    """
    A matrix initialized to be an invertible, orthogonal matrix, and is
    guarnteed to keep invertible during training.
    """

    def __init__(self,
                 seed_matrix: np.ndarray,
                 dtype: str = settings.float_x,
                 device: Optional[str] = None,
                 epsilon: float = EPSILON):
        """
        Construct a new :class:`StrictInvertibleMatrix`.

        Args:
            seed_matrix: A matrix that is used as a seed to obtain the
                initial invertible and orthogonal matrix.
            dtype: The dtype of the matrix.
            device: The device where to place new tensors and variables.
            epsilon: The infinitesimal constant to avoid dividing by zero or
                taking logarithm of zero.
        """
        initial_matrix = la.qr(seed_matrix)[0]
        device = device or current_device()

        super().__init__(initial_matrix.shape[0])
        matrix_shape = list(initial_matrix.shape)
        self.size = matrix_shape[0]

        initial_P, initial_L, initial_U = la.lu(initial_matrix)
        initial_s = np.diag(initial_U)
        initial_sign = np.sign(initial_s)
        initial_log_s = np.log(np.maximum(np.abs(initial_s), epsilon))
        initial_U = np.triu(initial_U, k=1)

        add_buffer(self, 'P', from_numpy(initial_P, dtype=dtype, device=device))
        assert_finite(
            add_parameter(
                self, 'pre_L', from_numpy(initial_L, dtype=dtype, device=device)),
            'pre_L',
        )
        add_buffer(
            self, 'L_mask', from_numpy(
                np.tril(np.ones(matrix_shape), k=-1), dtype=dtype, device=device)
        )
        assert_finite(
            add_parameter(self, 'pre_U', from_numpy(
                initial_U, dtype=dtype, device=device)),
            'pre_U',
        )
        add_buffer(
            self, 'U_mask', from_numpy(
                np.triu(np.ones(matrix_shape), k=1), dtype=dtype, device=device))
        add_buffer(
            self, 'sign', from_numpy(initial_sign, dtype=dtype, device=device))
        assert_finite(
            add_parameter(self, 'log_s', from_numpy(
                initial_log_s, dtype=dtype, device=device)),
            'log_s',
        )

    def forward(self,
                inverse: bool,
                compute_log_det: bool
                ) -> Tuple[Tensor, Optional[Tensor]]:
        P = self.P
        L = (self.L_mask * self.pre_L +
             torch.eye(self.size, dtype=P.dtype, device=self.P.device))
        U = self.U_mask * self.pre_U + torch.diag(self.sign * exp(self.log_s))

        log_det: Optional[Tensor] = None
        if inverse:
            matrix = matmul(
                matrix_inverse(U),
                matmul(matrix_inverse(L), matrix_inverse(P))
            )
            matrix = self._maybe_assert_finite(matrix, 'matrix', inverse)
            if compute_log_det:
                log_det = self._maybe_assert_finite(
                    -reduce_sum(self.log_s), 'log_det', inverse)
        else:
            matrix = matmul(P, matmul(L, U))
            matrix = self._maybe_assert_finite(matrix, 'matrix', inverse)
            if compute_log_det:
                log_det = self._maybe_assert_finite(
                    reduce_sum(self.log_s), 'log_det', inverse)

        return matrix, log_det

class InvertibleLinearNd(FeatureMappingFlow):
    """Base class for invertible linear transformation flows."""

    __constants__ = FeatureMappingFlow.__constants__ + (
        'num_features', 'strict', 'epsilon',
    )

    invertible_matrix: Module
    num_features: int
    strict: bool
    epsilon: float

    def __init__(self,
                 num_features: int,
                 strict: bool = False,
                 weight_init: TensorInitArgType = init.kaming_uniform,
                 dtype: str = settings.float_x,
                 device: Optional[str] = None,
                 epsilon: float = EPSILON):
        """
        Construct a new linear transformation flow.

        Args:
            num_features: The number of features to be transformed.
                The invertible transformation matrix will have the shape
                ``[num_features, num_features]``.
            strict: Whether or not to use the strict invertible matrix?
                Defaults to :obj:`False`.  See :class:`LooseInvertibleMatrix`
                and :class:`StrictInvertibleMatrix`.
            weight_init: The weight initializer for the seed matrix.
            dtype: The dtype of the invertible matrix.
            device: The device where to place new tensors and variables.
            epsilon: The infinitesimal constant to avoid having numerical issues.
        """
        spatial_ndims = self._get_spatial_ndims()
        device = device or current_device()

        super().__init__(
            axis=-(spatial_ndims + 1),
            event_ndims=(spatial_ndims + 1),
            explicitly_invertible=True,
        )

        self.num_features = int(num_features)
        self.strict = bool(strict)
        self.epsilon = float(epsilon)

        # Using the backend random generator instead of numpy generator
        # will allow the backend random seed to have effect on the initialization
        # step of the invertible matrix.
        seed_matrix = variable(
            shape=[num_features, num_features], dtype=dtype, device='cpu',
            initializer=weight_init, requires_grad=False,
        )
        seed_matrix = to_numpy(seed_matrix)

        if strict:
            self.invertible_matrix = StrictInvertibleMatrix(
                seed_matrix, dtype=dtype, device=device, epsilon=epsilon)
        else:
            self.invertible_matrix = LooseInvertibleMatrix(
                seed_matrix, dtype=dtype, device=device)

    def _get_spatial_ndims(self) -> int:
        raise NotImplementedError()

    def _affine_transform(self, input: Tensor, weight: Tensor) -> Tensor:
        raise NotImplementedError()

    @jit_method
    def _transform(self,
                   input: Tensor,
                   input_log_det: Optional[Tensor],
                   inverse: bool,
                   compute_log_det: bool
                   ) -> Tuple[Tensor, Optional[Tensor]]:
        # obtain the weight
        weight, log_det = self.invertible_matrix(
            inverse=inverse, compute_log_det=compute_log_det)
        spatial_ndims = self.x_event_ndims - 1
        weight = torch.reshape(weight, weight.shape + (1,) * spatial_ndims)

        # compute the output
        output = self._affine_transform(input, weight)

        # compute the log_det
        output_log_det = input_log_det
        if log_det is not None:
            log_det = log_det * torch.prod(
                torch.as_tensor(input.shape[input.dim() - spatial_ndims:],
                                dtype=log_det.dtype, device=log_det.device))
            if input_log_det is not None:
                output_log_det = input_log_det + log_det
            else:
                output_log_det = log_det

        return output, output_log_det

class InvertibleDense(InvertibleLinearNd):
    """An invertible linear transformation."""

    def _get_spatial_ndims(self) -> int:
        return 0

    @jit_method
    def _affine_transform(self, input: Tensor, weight: Tensor) -> Tensor:
        return torch.nn.functional.linear(input, weight)

class InvertibleConv1d(InvertibleLinearNd):
    """An invertible 1d 1x1 convolutional transformation."""

    def _get_spatial_ndims(self) -> int:
        return 1

    @jit_method
    def _affine_transform(self, input: Tensor, weight: Tensor) -> Tensor:
        return torch.nn.functional.conv1d(input, weight)

class InvertibleConv2d(InvertibleLinearNd):
    """An invertible 2d 1x1 convolutional transformation."""

    def _get_spatial_ndims(self) -> int:
        return 2

    @jit_method
    def _affine_transform(self, input: Tensor, weight: Tensor) -> Tensor:
        return torch.nn.functional.conv2d(input, weight)

class InvertibleConv3d(InvertibleLinearNd):
    """An invertible 3d 1x1 convolutional transformation."""

    def _get_spatial_ndims(self) -> int:
        return 3

    @jit_method
    def _affine_transform(self, input: Tensor, weight: Tensor) -> Tensor:
        return torch.nn.functional.conv3d(input, weight)

# ---- scale modules, for transforming input to output by a scale ----
class Scale(BaseValidateTensorLayer):
    """Base class for scaling `input`."""

    def _scale_and_log_scale(self,
                             pre_scale: Tensor,
                             inverse: bool,
                             compute_log_scale: bool
                             ) -> Tuple[Tensor, Optional[Tensor]]:
        raise NotImplementedError()

    def forward(self,
                input: Tensor,
                pre_scale: Tensor,
                event_ndims: int,
                input_log_det: Optional[Tensor] = None,
                compute_log_det: bool = True,
                inverse: bool = False
                ) -> Tuple[Tensor, Optional[Tensor]]:
        # validate the argument
        if input.dim() < event_ndims:
            raise ValueError(
                '`rank(input) >= event_ndims` does not hold: the `input` shape '
                'is {}, while `event_ndims` is {}.'.
                format(shape(input), event_ndims)
            )
        if pre_scale.dim() > input.dim():
            raise ValueError(
                '`rank(input) >= rank(pre_scale)` does not hold: the `input` '
                'shape is {}, while the shape of `pre_scale` is {}.'.
                format(shape(input), shape(pre_scale))
            )

        input_shape = input.shape
        event_ndims_start = len(input_shape) - event_ndims
        event_shape = input_shape[event_ndims_start:]
        log_det_shape = input_shape[: event_ndims_start]

        if input_log_det is not None:
            if input_log_det.shape != log_det_shape:
                raise ValueError(
                    'The shape of `input_log_det` is not expected: '
                    'expected to be {}, but got {}'.
                    format(list(log_det_shape), shape(input_log_det))
                )

        scale, log_scale = self._scale_and_log_scale(
            pre_scale, inverse, compute_log_det)

        if log_scale is not None:
            # the last `event_ndims` dimensions must match the `event_shape`
            r = log_scale.dim()
            if r < event_ndims or log_scale.shape[r - event_ndims:] != event_shape:
                # Note: equivalent as the following two lines, but compiles much slower
                #       on PyTorch 1.3.1 with JIT engine.
                # log_scale = broadcast_to_shape(
                #     log_scale,
                #     get_broadcast_shape(shape(log_scale), event_shape)
                # )
                log_scale = log_scale + torch.zeros(
                    event_shape, device=log_scale.device, dtype=log_scale.dtype)

                r = log_scale.dim()
                if log_scale.shape[r - event_ndims:] != event_shape:
                    raise ValueError(
                        'The shape of the final {}d of `log_scale` is not '
                        'expected: expected to be {}, but got {}.'.
                        format(event_ndims, event_shape, log_scale.shape[r - event_ndims:])
                    )

            # reduce the last `event_ndims` of log_scale
            log_scale = reduce_sum(log_scale, axis=int_range(-event_ndims, 0))

            # now add to input_log_det, or broadcast `log_scale` to `log_det_shape`
            if input_log_det is not None:
                output_log_det = input_log_det + log_scale
            else:
                output_log_det = log_scale
                if output_log_det.shape != log_det_shape:
                    output_log_det = output_log_det + torch.zeros(
                        log_det_shape, device=output_log_det.device,
                        dtype=output_log_det.dtype
                    )

            if output_log_det.shape != log_det_shape:
                raise ValueError(
                    'The shape of the computed `output_log_det` is not expected: '
                    'expected to be {}, but got {}.'.
                    format(shape(output_log_det), list(log_det_shape))
                )
        else:
            output_log_det = None

        return input * scale, output_log_det

class ExpScale(Scale):
    """
    Scaling `input` with `exp` activation.

    ::

        if inverse:
            output = input / exp(pre_scale)
            output_log_det = -pre_scale
        else:
            output = input * exp(pre_scale)
            output_log_det = pre_scale
    """

    def _scale_and_log_scale(self,
                             pre_scale: Tensor,
                             inverse: bool,
                             compute_log_scale: bool
                             ) -> Tuple[Tensor, Optional[Tensor]]:
        log_scale: Optional[Tensor] = None

        the_pre_scale = -pre_scale if inverse else pre_scale
        scale = self._maybe_assert_finite(exp(the_pre_scale), 'scale', inverse)
        if compute_log_scale:
            log_scale = self._maybe_assert_finite(
                the_pre_scale, 'log_scale', inverse)

        return scale, log_scale

class SigmoidScale(Scale):
    """
    Scaling `input` with `sigmoid` activation.

    ::

        if inverse:
            output = input / sigmoid(pre_scale)
            output_log_det = -log(sigmoid(pre_scale))
        else:
            output = input * sigmoid(pre_scale)
            output_log_det = log(sigmoid(pre_scale))
    """

    __constants__ = Scale.__constants__ + ('pre_scale_bias',)

    pre_scale_bias: float

    def __init__(self, pre_scale_bias: float = 0.):
        super().__init__()
        self.pre_scale_bias = pre_scale_bias

    def _scale_and_log_scale(self,
                             pre_scale: Tensor,
                             inverse: bool,
                             compute_log_scale: bool
                             ) -> Tuple[Tensor, Optional[Tensor]]:
        if self.pre_scale_bias != 0.:
            pre_scale = pre_scale + self.pre_scale_bias

        log_scale: Optional[Tensor] = None
        if inverse:
            neg_pre_scale = -pre_scale
            scale = self._maybe_assert_finite(
                exp(neg_pre_scale) + 1., 'scale', inverse)
            if compute_log_scale:
                log_scale = self._maybe_assert_finite(
                    softplus(neg_pre_scale), 'log_scale', inverse)
        else:
            scale = self._maybe_assert_finite(sigmoid(pre_scale), 'scale', inverse)
            if compute_log_scale:
                log_scale = self._maybe_assert_finite(
                    -softplus(-pre_scale), 'log_scale', inverse)

        return scale, log_scale

class LinearScale(Scale):
    """
    Scaling `input` with `linear` activation.

    ::

        if inverse:
            output = input / pre_scale
            output_log_det = -log(abs(pre_scale))
        else:
            output = input * pre_scale
            output_log_det = log(abs(pre_scale))
    """

    __constants__ = Scale.__constants__ + ('epsilon',)

    epsilon: float

    def __init__(self, epsilon: float = EPSILON):
        super().__init__()
        self.epsilon = epsilon

    def _scale_and_log_scale(self,
                             pre_scale: Tensor,
                             inverse: bool,
                             compute_log_scale: bool
                             ) -> Tuple[Tensor, Optional[Tensor]]:
        log_scale: Optional[Tensor] = None
        epsilon = float_scalar_like(self.epsilon, pre_scale)

        if inverse:
            scale = self._maybe_assert_finite(1. / pre_scale, 'scale', inverse)
            if compute_log_scale:
                log_scale = self._maybe_assert_finite(
                    -log(maximum(abs(pre_scale), epsilon)), 'log_scale', inverse)
        else:
            scale = self._maybe_assert_finite(pre_scale, 'scale', inverse)
            if compute_log_scale:
                log_scale = self._maybe_assert_finite(
                    log(maximum(abs(pre_scale), epsilon)), 'log_scale', inverse)

        return scale, log_scale