# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function, unicode_literals
import torch
import copy
import six  # NOQA
import torch.nn as nn
import torchvision
import ubelt as ub
import numpy as np
from collections import OrderedDict
from netharn.analytic.output_shape_for import OutputShapeFor
from netharn.analytic import analytic_for
from distutils.version import LooseVersion
# try:
# from netharn.device import MountedModel
# except ImportError:
#     MountedModel = None

REGISTERED_TYPES = []

if LooseVersion(torch.__version__) >= LooseVersion('1.5.0'):
    CONV_TRANSPOSE_TYPES = (nn.modules.conv._ConvTransposeNd,)
else:
    CONV_TRANSPOSE_TYPES = (nn.modules.conv._ConvTransposeMixin,)

def ensure_array_nd(data, n):
    if ub.iterable(data):
        return np.array(data)
    else:
        return np.array([data] * n)

def compute_type(*types):
    def _wrap(func):
        for type in types:
            if type is not None:
                REGISTERED_TYPES.append((type, func))
        return func
    return _wrap

class ReceptiveFieldTypeError(TypeError):
    pass

class ReceptiveField(OrderedDict, analytic_for.Output):
    """
    container for holding a receptive feild

    Example:
        >>> self = ReceptiveField.coerce({
        >>>     'stride': np.array([4]),
        >>>     'shape': np.array([1]),
        >>>     'crop': np.array([0]),
        >>> })
        >>> self_copy = copy.deepcopy(self)
    """
    def __init__(self, data, hidden=None):
        # Inheriting from an odict consistently between python 2/3 is weird
        data2 = OrderedDict(sorted(OrderedDict(data).items()))
        OrderedDict.__init__(self, data2)
        self.data = data2
        self.hidden = hidden

    def __copy__(self):
        self_copy = ReceptiveField(self.data, self.hidden)
        return self_copy

    def __deepcopy__(self, memo):
        data_copy = copy.deepcopy(self.data, memo)
        hidden_copy = copy.deepcopy(self.hidden, memo)
        self_copy = ReceptiveField(data_copy, hidden_copy)
        return self_copy

    @classmethod
    def coerce(cls, data, hidden=None):
        """
        Example:
            >>> # test weird python2 failure case
            >>> from netharn.analytic.receptive_field_for import *
            >>> cls = ReceptiveField
            >>> data = [(0, ReceptiveFieldFor.input())]
            >>> self = cls.coerce(data)
            >>> print(ub.repr2(self, with_dtype=False))
            {
                0: {
                    'crop': np.array([0., 0.]),
                    'shape': np.array([1., 1.]),
                    'stride': np.array([1., 1.]),
                },
            }
        """
        # TODO: make this work like OutputShape
        if data is None:
            self = ReceptiveFieldFor.input()
            self.hidden = hidden
        elif isinstance(data, cls):
            if hidden is None:
                self = data
            else:
                self = data.__class__(data, hidden)
        else:
            self = cls(data, hidden)
        return self

    # def __getitem__(self, key):
    #     return self.data[key]

class HiddenFields(analytic_for.Hidden):
    """
    Augments normal hidden fields dicts with a convinience setitem
    """
    pass

class _TorchMixin(object):
    """
    Receptive field formulas for PyTorch primatives
    """

    @staticmethod
    def input(input_field=None, n=2):
        """
        Basic input receptive field is just a single pixel.
        """
        if input_field is not None:
            raise ValueError('nothing can precede the input')
        input_field = ReceptiveField.coerce({
            # The input receptive field stride / scale factor is 1.
            'stride': ensure_array_nd(1.0, n),
            # The input receptive field shape is 1 pixel.
            'shape': ensure_array_nd(1.0, n),
            # Use the coordinate system where the top left corner is 0, 0 ( This is unlike [1], which uses 0.5)
            'crop': ensure_array_nd(0.0, n),
        })
        return input_field

    @staticmethod
    def _kernelized(module, input_field=None, ndim=None):
        """
        Receptive field formula for general sliding kernel based layers
        This works for both convolutional and pooling layers.

        Notes:
            Baseline formulas are from [1]. Information about how to include
            dilation (atrous) convolutions can be found in [2, 3].  Better info
            seems to be available in [4].

            * tensorflow has similar functionality
            https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/receptive_field/python/util/receptive_field.py

            * To preserve spatial extent, padding should equal `(k - 1) * d / 2`.

        References:
            [1] https://medium.com/mlreview/a-guide-to-receptive-field-arithmetic-for-convolutional-neural-networks-e0f514068807
            [2] http://www.erogol.com/dilated-convolution/
            [3] https://stackoverflow.com/questions/35582521/how-to-calculate-receptive-field-shape
            [4] https://arxiv.org/pdf/1603.07285.pdf

        Example:
            >>> module = nn.Conv2d(1, 1, kernel_size=5, stride=2, padding=2, dilation=3)
            >>> field = ReceptiveFieldFor._kernelized(module)
            >>> print(ub.repr2(field, nl=0, with_dtype=False))
            {'crop': np.array([4., 4.]), 'shape': np.array([13., 13.]), 'stride': np.array([2., 2.])}

            >>> module = nn.MaxPool2d(kernel_size=3, stride=2, padding=2, dilation=2)
            >>> field = ReceptiveFieldFor._kernelized(module)
            >>> print(ub.repr2(field, nl=0, with_dtype=False))
            {'crop': np.array([0., 0.]), 'shape': np.array([5., 5.]), 'stride': np.array([2., 2.])}

            >>> module = nn.MaxPool2d(kernel_size=3, stride=2, padding=2, dilation=1)
            >>> field = ReceptiveFieldFor._kernelized(module)
            >>> print(ub.repr2(field, nl=0, with_dtype=False))
            {'crop': np.array([-1., -1.]), 'shape': np.array([3., 3.]), 'stride': np.array([2., 2.])}

            >>> module = nn.AvgPool2d(kernel_size=3, stride=2, padding=2)
            >>> field = ReceptiveFieldFor._kernelized(module)
            >>> print(ub.repr2(field, nl=0, with_dtype=False))
            {'crop': np.array([-1., -1.]), 'shape': np.array([3., 3.]), 'stride': np.array([2., 2.])}
        """
        # impl = ReceptiveFieldFor.impl
        if input_field is None:
            input_field = ReceptiveFieldFor.input()

        # Hack to get the number of space-time dimensions
        if ndim is None:
            try:
                if module.__class__.__name__.endswith('1d'):
                    ndim = 1
                elif module.__class__.__name__.endswith('2d'):
                    ndim = 2
                elif module.__class__.__name__.endswith('3d'):
                    ndim = 3
            except AttributeError:
                if module.__name__.endswith('1d'):
                    ndim = 1
                elif module.__name__.endswith('2d'):
                    ndim = 2
                elif module.__name__.endswith('3d'):
                    ndim = 3
        if ndim is None:
            raise ValueError('Cannot infer ndim from {}'.format(module))

        k = ensure_array_nd(module.kernel_size, ndim)
        s = ensure_array_nd(module.stride, ndim)
        p = ensure_array_nd(module.padding, ndim)
        d = ensure_array_nd(getattr(module, 'dilation', 1), ndim)

        # To calculate receptive feild we first need to find the SUPPORT of
        # this layer. The support is the number/extent of extra surrounding
        # pixels adding this layer will take into account. Given this, we can
        # compute the receptive feild wrt the original input by combining this
        # information with the previous receptive feild.
        #
        # In the normal case (with no dilation, d=1) the support is (k - 1).
        # This is because because the operation is able to see a window of shape
        # k in the input, and produces a single output pixel (hence the k). The
        # center input pixel corresponds with the output, so it does not expand
        # the receptive feild (hence the -1), but all other input pixels do
        # expand the field (thus the k-1).
        #
        # The stride of this layer will not affect the support.
        #
        # The dilation of the current layer DOES impact the support.
        # This expands the effective kernel shape, but it does cause the data
        # each operation sees to become more diffuse. However, even though what
        # it sees in that extent is more diffuse, the RF is just a bound, so we
        # can ignore the diffuseness effect and simply scale the input kernel
        # shape by the dilation amount. Hense we get
        support = (k - 1) * d

        """
        Note the above is correct because:

            import sympy as sym
            k, d = sym.symbols('k, d')

            # Compute the support from formula in 5.1 of [4]
            # To understand the relationship tying the dilation rate d and the
            # output shape o, it is useful to think of the impact of d on the
            # effective kernel shape. A kernel of shape k dilated by a factor d
            # has an effective shape.
            effective_kernel_size = k + (k - 1) * (d - 1)
            support_v1 = sym.expand(effective_kernel_size - 1)

            # Compute support from our method
            support_v2 = sym.expand((k - 1) * d)

            # They are equivalent. QED
            assert sym.Eq(support_v1, support_v2)
        """

        # Compute how many pixels this layer takes off the side Note that an
        # even shape kernel results in half pixel crops.  This is expected and
        # correct. To use the crop in practice take the floor / ceil of the
        # final result, but in this intermediate stage, subpixel crops are
        # perfectly valid.
        crop = ((support / 2.0) - p)

        field = ReceptiveField.coerce({
            # The new stride only depends on the layer stride and the previous
            # stride.
            'stride': input_field['stride'] * s,

            # The stride of the current layer does not impact the receptive
            # feild, however the stride of the previous layer does. This is
            # because each pixel in the incoming layer really corresponds
            # `input_field['stride']` pixels in the original input.
            'shape':   input_field['shape'] + support * input_field['stride'],

            # Padding does not influence the RF shape, but it does influence
            # where the start pixel is (i.e. without the right amount of
            # padding the the edge of the previous layer is cropped).
            'crop': input_field['crop'] + crop * input_field['stride'],
        })
        return field

    @staticmethod
    def _unchanged(module, input_field=None):
        """ Formula for layers that do not change the receptive field """
        if input_field is None:
            input_field = ReceptiveFieldFor.input()
        return input_field

    @staticmethod
    @compute_type(nn.Linear)
    def linear(module, input_field=None):
        # Linear layers (sort-of) dont change the RF
        return ReceptiveFieldFor._unchanged(module, input_field)
        # Perhaps we could do this if we knew the input shape
        # raise NotImplementedError(
        #     'Cannot compute receptive field shape on a Linear layer')

    @staticmethod
    def _kernelized_tranpose(module, input_field=None):
        """
        Receptive field formula for pooling layers

        Example:
            >>> from netharn.analytic.receptive_field_for import *
            >>> from netharn.analytic.output_shape_for import *
            >>> module = nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2, padding=2)
            >>> ReceptiveFieldFor(module)()

            >>> # This network should effectively invert itself
            >>> module = nn.Sequential(ub.odict([
            >>>     #('a', nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1)),
            >>>     ('c1', nn.Conv2d(1, 1, kernel_size=3, stride=2)),
            >>>     ('c2', nn.Conv2d(1, 1, kernel_size=3, stride=2)),
            >>>     ('c3', nn.Conv2d(1, 1, kernel_size=3, stride=2)),
            >>>     ('c3T', nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2)),
            >>>     ('c2T', nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2)),
            >>>     ('c1T', nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2)),
            >>> ]))
            >>> print(ub.repr2(ReceptiveFieldFor(module)()))
            >>> ReceptiveFieldFor(module)()
            >>> OutputShapeFor(module)._check_consistency([1, 1, 32, 32])

            >>> module = nn.Sequential(ub.odict([
            >>>     #('a', nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1)),
            >>>     ('c1', nn.Conv2d(1, 1, kernel_size=3, stride=2, dilation=2)),
            >>>     ('c2', nn.Conv2d(1, 1, kernel_size=3, stride=2, dilation=2)),
            >>>     ('c3', nn.Conv2d(1, 1, kernel_size=3, stride=2, dilation=2)),
            >>>     ('c3T', nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2, dilation=2)),
            >>>     ('c2T', nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2, dilation=2)),
            >>>     ('c1T', nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2, dilation=2)),
            >>> ]))
            >>> print(ub.repr2(ReceptiveFieldFor(module)()))

            >>> # This network is pathological
            >>> module = nn.Sequential(ub.odict([
            >>>     #('a', nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1)),
            >>>     ('c1', nn.Conv2d(1, 1, kernel_size=3, stride=7, dilation=2)),
            >>>     ('c2', nn.Conv2d(1, 1, kernel_size=5, stride=6, padding=1)),
            >>>     ('c3', nn.Conv2d(1, 1, kernel_size=7, stride=5)),
            >>>     ('c3T', nn.ConvTranspose2d(1, 1, kernel_size=7, stride=6)),
            >>>     ('c2T', nn.ConvTranspose2d(1, 1, kernel_size=5, stride=7, padding=1)),
            >>>     ('c1T', nn.ConvTranspose2d(1, 1, kernel_size=3, stride=8, dilation=2)),
            >>> ]))
            >>> print(ub.repr2(ReceptiveFieldFor(module)()))
            >>> ReceptiveFieldFor(module)()
            >>> OutputShapeFor(module)([1, 1, 900, 900])
            >>> OutputShapeFor(module)([1, 1, 900, 900]).hidden
            >>> OutputShapeFor(module)._check_consistency([1, 1, 900, 900])

            >>> module = nn.Sequential(
            >>>     nn.Conv2d(1, 1, kernel_size=3, stride=2),
            >>>     nn.Conv2d(1, 1, kernel_size=3, stride=2),
            >>>     nn.Conv2d(1, 1, kernel_size=3, stride=2),
            >>>     nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2),
            >>>     nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2),
            >>>     nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2),
            >>> )
            >>> ReceptiveFieldFor(module)()

            >>> module = nn.Conv2d(1, 1, kernel_size=3, stride=2, padding=1)
            >>> ReceptiveFieldFor(module)()

            >>> OutputShapeFor(nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2, padding=0, output_padding=(1, 1)))._check_consistency([1, 1, 1, 1])

            >>> # Figure 4.4
            >>> OutputShapeFor(nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=2))([1, 1, 5, 5])
            >>> OutputShapeFor(nn.ConvTranspose2d(1, 1, kernel_size=3, stride=1, padding=2))._check_consistency([1, 1, 5, 5])
            >>> OutputShapeFor(nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=0))([1, 1, 7, 7])

            >>> # Figure 4.5
            >>> OutputShapeFor(nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2, padding=0))._check_consistency([1, 1, 5, 5])
            >>> OutputShapeFor(nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=0))([1, 1, 7, 7])

            >>> ReceptiveFieldFor(module)()
        """
        # impl = ReceptiveFieldFor.impl
        if input_field is None:
            input_field = ReceptiveFieldFor.input()

        # Hack to get the number of space-time dimensions
        ndim = None
        try:
            if module.__name__.endswith('1d'):
                ndim = 1
            elif module.__name__.endswith('2d'):
                ndim = 2
            elif module.__name__.endswith('3d'):
                ndim = 3
        except AttributeError:
            pass

        if ndim is None:
            if hasattr(module, '_dim'):
                ndim = module._dim

        # A non-trivial transpose convolution should:
        # * decrease the stride (because the stride is fractional)
        # the padding has to be equal to the shape of the kernel minus one
        """
        From [4]:

        A convolution described by k, s and p has an associated transposed convolution described by:
        * k' = k,
        * s' = 1,
        * p' = k - p - 1,
        * i' = the shape of the stretched input obtained by adding s − 1 zeros
            between each input unit,
        * a = (i + 2p − k) % s, represents the number of zeros added to the
         bottom and right edges of the input,

         And has output shape:
             o' = s(i' - 1) + a + k - 2p

        For convT it is always the case that s'=1, howver, note that s' is not
        what we use to compute the new stride of the output, because that is
        actually a fractional stride.
        """

        # Definitions:
        # In the following comments we discuss 3 distinct layers
        # (1) The original convolution (conv)
        # (2) The transpose convolution that inverts the original (convT)
        # (3) The regular convolution that is equivalent to the transpose
        # convolution given a specially transformed input tensor (convE)

        # The parameters of a convT are actually the parameters of conv, the
        # convolution we are trying to "undo", but we will refer to them as
        # parameters of convT (because they are that as well).
        k_ = ensure_array_nd(module.kernel_size, ndim)
        s_ = ensure_array_nd(module.stride, ndim)
        p_ = ensure_array_nd(module.padding, ndim)
        d_ = ensure_array_nd(getattr(module, 'dilation', 1), ndim)

        # TODO: incorporate output padding and right-side padding / cropping
        # Note: output padding does not impact the receptive field, however it
        # does cause some "right-side" croping, which we are not computing here
        # yet.

        out_pad = ensure_array_nd(module.output_padding, ndim)  # NOQA
        # if not np.all(out_pad == 0):
        #     raise NotImplementedError('cannot handle nonzero output_padding yet')

        # Howver, there is an equivalent way of forumulating a convT as convE:
        # a regular conv applied on a specially padded input tensor.
        # The parameters that define convE are:
        k = k_
        d = d_
        s = 1  # stride is always 1 because of the special input transform
        # p = k_ - p_ - 1  # NOTE: original formula likely assumed dilation=1
        p = (k_ - 1) * d_ - p_

        # In order for convE to be equivalent to convT, we need to apply convE
        # to a specially transformed (padded) input tensor.
        # The padding applied to the input tensor puts extra zeros between each
        # row/col. The number of extra zeros is the stride of the convT - 1.
        # The left and right sides of the input tensor are also padded but that
        # wont factor into the RF calculation.
        extra_zeros = s_ - 1
        # This means that the effective support added to the RF shape by convE
        # will be less than it normally would because we don't count the extra
        # zeros in our transformed input as real pixels.
        effective_support = (k - 1 - extra_zeros) * d
        # NOTE; if the stride is larger than the kernel, some output pixels
        # will actually just be zeros and have no receptive feild.
        effective_support = np.maximum(0, effective_support)

        # This special input transform also has the effect of decreasing the RF
        # stride.  Transpose conv are sometimes called fractional-stride
        # convolutions This is because they have an effective stride of 1 / s_
        effective_stride = 1 / s_

        # We calculate the support of convE as if were applied to a normal
        # input tensor in order to calculate how the start (top-left) pixel
        # position is modified.
        support = (k - 1) * d

        # After transformation the effective stride of the input is
        effective_input_stride = input_field['stride'] * effective_stride

        # how many pixels does this layer crop off the sides of the input
        crop = ((support / 2) - p)

        # print('effective_support = {!r}'.format(effective_support))

        field = ReceptiveField.coerce({
            # The new stride only depends on the layer stride and the previous
            # stride.
            'stride': effective_input_stride * s,

            # The stride of the current layer does not impact the receptive
            # feild, however the stride of the previous layer does. This is
            # because each pixel in the incoming layer really corresponds
            # `input_field['stride']` pixels in the original input.
            'shape':   input_field['shape'] + effective_support * input_field['stride'],

            # Padding does not influence the RF shape, but it does influence
            # where the start pixel is (i.e. without the right amount of
            # padding the the edge of the previous layer is cropped).
            'crop': input_field['crop'] + crop * effective_input_stride,
        })

        return field
        # raise NotImplementedError('todo')

    @compute_type(*CONV_TRANSPOSE_TYPES)
    def convT(module, input_field=None):
        return ReceptiveFieldFor._kernelized_tranpose(module, input_field)

    @compute_type(nn.modules.conv.Conv1d, nn.modules.conv.Conv2d, nn.modules.conv.Conv3d)
    def convnd(module, input_field=None):
        return ReceptiveFieldFor._kernelized(module, input_field)

    @staticmethod
    @compute_type(nn.modules.pooling._MaxPoolNd)
    def maxpoolnd(module, input_field=None):
        return ReceptiveFieldFor._kernelized(module, input_field)

    @staticmethod
    @compute_type(nn.modules.pooling._AvgPoolNd)
    def avepoolnd(module, input_field=None):
        return ReceptiveFieldFor._kernelized(module, input_field)

    @staticmethod
    @compute_type(nn.modules.pooling._AdaptiveMaxPoolNd, nn.modules.pooling._AdaptiveAvgPoolNd)
    def adaptive_avepoolnd(module, input_field=None):
        """
        it is not possible to analytically compute an adaptive receptive field.

        References:
            https://forums.fast.ai/t/ideas-behind-adaptive-max-pooling/12634/3
            https://arxiv.org/abs/1406.4729
        """
        raise Exception('not possible to compute adaptive RF without knowning the input_shape ahead of time')
        # return ReceptiveFieldFor._kernelized(module, input_field)

    @staticmethod
    @compute_type(nn.ReLU)
    def relu(module, input_field=None):
        return ReceptiveFieldFor._unchanged(module, input_field)

    @staticmethod
    @compute_type(nn.ReLU6, nn.PReLU, nn.LeakyReLU, nn.ELU, nn.CELU, nn.SELU)
    def _unchanged_activation(module, input_field=None):
        return ReceptiveFieldFor._unchanged(module, input_field)

    @staticmethod
    @compute_type(nn.functional.relu, nn.functional.relu6)
    def _unchanged_activation_func(input_field=None):
        # return ReceptiveFieldFor._unchanged(module, input_field)
        return ReceptiveFieldFor._unchanged(None, input_field)

    @staticmethod
    @compute_type(nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d,
                  nn.modules.normalization.GroupNorm,
                  nn.modules.normalization.LocalResponseNorm,
                  nn.modules.normalization.LayerNorm, nn.CrossMapLRN2d,
                  nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d)
    def normalization(module, input_field=None):
        return ReceptiveFieldFor._unchanged(module, input_field)

    @staticmethod
    @compute_type(nn.modules.dropout._DropoutNd)
    def dropout(module, input_field=None):
        return ReceptiveFieldFor._unchanged(module, input_field)

    @staticmethod
    @compute_type(nn.Sequential)
    def sequential(module, input_field=None):
        """
        Example:
            >>> import netharn as nh
            >>> self = nn.Sequential(
            >>>     nn.Conv2d(2, 3, kernel_size=3),
            >>>     nn.Conv2d(3, 5, kernel_size=3),
            >>>     nn.Conv2d(5, 7, kernel_size=3),
            >>> )
            >>> rfield = nh.ReceptiveFieldFor(self)()
            >>> print('rfield = {}'.format(ub.repr2(rfield, nl=1, with_dtype=False)))
            rfield = {
                'crop': np.array([3., 3.]),
                'shape': np.array([7., 7.]),
                'stride': np.array([1., 1.]),
            }
        """
        if input_field is None:
            input_field = ReceptiveFieldFor.input()
        rfield = input_field
        hidden = HiddenFields()
        iter_ = iter(module._modules.items())
        for key, child in module._modules.items():
            key, child = next(iter_)
            if hasattr(child, 'receptive_field_for'):
                rfield = hidden[key] = child.receptive_field_for(rfield)
            else:
                rfield = hidden[key] = ReceptiveFieldFor(child)(rfield)
        rfield = ReceptiveField.coerce(rfield)
        rfield.hidden = hidden
        return rfield

    @staticmethod
    @compute_type(torch.nn.DataParallel)
    def data_parallel(module, *args, **kw):
        return ReceptiveFieldFor(module.module)(*args, **kw)

class _TorchvisionMixin(object):
    """
    Compute receptive fields for components of torchvision models
    """

    @staticmethod
    @compute_type(torchvision.models.resnet.BasicBlock)
    def resent_basic_block(module, input_field=None):
        """
        Example:
            >>> # xdoctest: +REQUIRES(--network)
            >>> import torchvision  # NOQA
            >>> module = torchvision.models.resnet18().layer1[0]
            >>> field = ReceptiveFieldFor(module)()
            >>> print(ub.repr2(field.hidden, nl=1, with_dtype=False))
            {
                'conv1': {'crop': np.array([0., 0.]), 'shape': np.array([3., 3.]), 'stride': np.array([1., 1.])},
                'bn1': {'crop': np.array([0., 0.]), 'shape': np.array([3., 3.]), 'stride': np.array([1., 1.])},
                'relu1': {'crop': np.array([0., 0.]), 'shape': np.array([3., 3.]), 'stride': np.array([1., 1.])},
                'conv2': {'crop': np.array([0., 0.]), 'shape': np.array([5., 5.]), 'stride': np.array([1., 1.])},
                'bn2': {'crop': np.array([0., 0.]), 'shape': np.array([5., 5.]), 'stride': np.array([1., 1.])},
                'relu2': {'crop': np.array([0., 0.]), 'shape': np.array([5., 5.]), 'stride': np.array([1., 1.])},
            }
        """
        if input_field is None:
            input_field = ReceptiveFieldFor.input()
        hidden = HiddenFields()

        rfield = input_field

        rfield = hidden['conv1'] = ReceptiveFieldFor(module.conv1)(rfield)
        rfield = hidden['bn1'] = ReceptiveFieldFor(module.bn1)(rfield)
        rfield = hidden['relu1'] = ReceptiveFieldFor(module.relu)(rfield)

        rfield = hidden['conv2'] = ReceptiveFieldFor(module.conv2)(rfield)
        rfield = hidden['bn2'] = ReceptiveFieldFor(module.bn2)(rfield)
        rfield = hidden['relu2'] = ReceptiveFieldFor(module.relu)(rfield)

        if module.downsample is not None:
            hidden['downsample'] = ReceptiveFieldFor(module.downsample)(input_field)

        rfield = ReceptiveFieldFor(module.relu)(rfield)
        rfield.hidden = hidden
        return rfield

    @staticmethod
    @compute_type(torchvision.models.resnet.Bottleneck)
    def resent_bottleneck(module, input_field=None):
        """
        CommandLine:
            xdoctest -m netharn.analytic.receptive_field_for _TorchvisionMixin.resent_bottleneck --network

        Example:
            >>> # xdoctest: +REQUIRES(--network)
            >>> import torchvision  # NOQA
            >>> module = torchvision.models.resnet50().layer1[0]
            >>> field = ReceptiveFieldFor(module)()
            >>> print(ub.repr2(field.hidden.shallow(1), nl=1, with_dtype=False))
            {
                'conv1': {'crop': ...([0., 0.]), 'shape': ...([1., 1.]), 'stride': ...([1., 1.])},
                'bn1': {'crop': ...([0., 0.]), 'shape': ...([1., 1.]), 'stride': ...([1., 1.])},
                'relu1': {'crop': ...([0., 0.]), 'shape': ...([1., 1.]), 'stride': ...([1., 1.])},
                'conv2': {'crop': ...([0., 0.]), 'shape': ...([3., 3.]), 'stride': ...([1., 1.])},
                'bn2': {'crop': ...([0., 0.]), 'shape': ...([3., 3.]), 'stride': ...([1., 1.])},
                'relu2': {'crop': ...([0., 0.]), 'shape': ...([3., 3.]), 'stride': ...([1., 1.])},
                'conv3': {'crop': ...([0., 0.]), 'shape': ...([3., 3.]), 'stride': ...([1., 1.])},
                'bn3': {'crop': ...([0., 0.]), 'shape': ...([3., 3.]), 'stride': ...([1., 1.])},
                'downsample': {'crop': ...([0., 0.]), 'shape': ...([1., 1.]), 'stride': ...([1., 1.])},
            }
        """
        if input_field is None:
            input_field = ReceptiveFieldFor.input()
        rfield = input_field
        hidden = HiddenFields()

        rfield = hidden['conv1'] = ReceptiveFieldFor(module.conv1)(rfield)
        rfield = hidden['bn1'] = ReceptiveFieldFor(module.bn1)(rfield)
        rfield = hidden['relu1'] = ReceptiveFieldFor(module.relu)(rfield)

        rfield = hidden['conv2'] = ReceptiveFieldFor(module.conv2)(rfield)
        rfield = hidden['bn2'] = ReceptiveFieldFor(module.bn2)(rfield)
        rfield = hidden['relu2'] = ReceptiveFieldFor(module.relu)(rfield)

        rfield = hidden['conv3'] = ReceptiveFieldFor(module.conv3)(rfield)
        rfield = hidden['bn3'] = ReceptiveFieldFor(module.bn3)(rfield)

        if module.downsample is not None:
            hidden['downsample'] = ReceptiveFieldFor(module.downsample)(input_field)

        rfield = ReceptiveFieldFor(module.relu)(rfield)
        rfield.hidden = hidden
        return rfield

    @staticmethod
    @compute_type(torchvision.models.resnet.ResNet)
    def resnet_model(module, input_field=None, input_shape=None):
        """
        CommandLine:
            xdoctest -m netharn.analytic.receptive_field_for _TorchvisionMixin.resnet_model --network

        Example:
            >>> # DISABLE_DOCTEST
            >>> # Note: newest torchvision breaks this
            >>> # xdoctest: +REQUIRES(--network)
            >>> from netharn.analytic.receptive_field_for import *
            >>> module = torchvision.models.resnet50()
            >>> input_shape = (1, 3, 224, 224)
            >>> field = ReceptiveFieldFor(module)(input_shape=input_shape)
            >>> print(ub.repr2(field.hidden.shallow(1), nl=1, with_dtype=False))
            {
                'conv1': {'crop': ...([0., 0.]), 'shape': ...([7., 7.]), 'stride': ...([2., 2.])},
                'bn1': {'crop': ...([0., 0.]), 'shape': ...([7., 7.]), 'stride': ...([2., 2.])},
                'relu1': {'crop': ...([0., 0.]), 'shape': ...([7., 7.]), 'stride': ...([2., 2.])},
                'maxpool': {'crop': ...([0., 0.]), 'shape': ...([11., 11.]), 'stride': ...([4., 4.])},
                'layer1': {'crop': ...([0., 0.]), 'shape': ...([35., 35.]), 'stride': ...([4., 4.])},
                'layer2': {'crop': ...([0., 0.]), 'shape': ...([91., 91.]), 'stride': ...([8., 8.])},
                'layer3': {'crop': ...([0., 0.]), 'shape': ...([267., 267.]), 'stride': ...([16., 16.])},
                'layer4': {'crop': ...([0., 0.]), 'shape': ...([427., 427.]), 'stride': ...([32., 32.])},
                'avgpool': {'crop': ...([96., 96.]), 'shape': ...([619., 619.]), 'stride': ...([32., 32.])},
                'flatten': {'crop': ...([96., 96.]), 'shape': ...([811., 811.]), 'stride': ...([32., 32.])},
                'fc': {'crop': ...([96., 96.]), 'shape': ...([811., 811.]), 'stride': ...([32., 32.])},
            }

        """
        if input_field is None:
            input_field = ReceptiveFieldFor.input()
        rfield = input_field
        hidden = HiddenFields()
        rfield = hidden['conv1'] = ReceptiveFieldFor(module.conv1)(rfield)
        rfield = hidden['bn1'] = ReceptiveFieldFor(module.bn1)(rfield)
        rfield = hidden['relu1'] = ReceptiveFieldFor(module.relu)(rfield)
        rfield = hidden['maxpool'] = ReceptiveFieldFor(module.maxpool)(rfield)

        rfield = hidden['layer1'] = ReceptiveFieldFor(module.layer1)(rfield)
        rfield = hidden['layer2'] = ReceptiveFieldFor(module.layer2)(rfield)
        rfield = hidden['layer3'] = ReceptiveFieldFor(module.layer3)(rfield)
        rfield = hidden['layer4'] = ReceptiveFieldFor(module.layer4)(rfield)

        rfield = hidden['avgpool'] = ReceptiveFieldFor(module.avgpool)(rfield)

        if input_shape is None:
            raise ValueError('input shape is required')

        output_shape = OutputShapeFor(module)(input_shape)
        avgpool_shape = output_shape.hidden.shallow(1)['layer4']
        spatial_shape = np.array(avgpool_shape[2:])

        # Keep everything the same except increase the RF shape
        # based on how many output pixels there are.
        rfield_flatten = ReceptiveField.coerce(dict(**rfield))
        # not sure if this is 100% correct
        rfield_flatten['shape'] = rfield['shape'] + (spatial_shape - 1) * rfield['stride']
        rfield = hidden['flatten'] = rfield_flatten

        # The reshape operation will blend the receptive fields of the inputs
        # but it will depend on the output shape of the layer.
        # rfield = (rfield[0], prod(rfield[1:]))

        rfield = hidden['fc'] = ReceptiveFieldFor(module.fc)(rfield)
        rfield.hidden = hidden
        return rfield

class ReceptiveFieldFor(analytic_for.OutputFor, _TorchMixin, _TorchvisionMixin):
    """
    Knows how to compute the receptive fields for many pytorch primatives and
    some torchvision components.

    References:
        https://medium.com/mlreview/a-guide-to-receptive-field-arithmetic-for-convolutional-neural-networks-e0f514068807

    Returns:
        Tuple[object, Dict]:
            fields: object: The hidden layer recepvive fields (can be complex due to nesting)
            field: Dict: a dictionary containing receptive field information.

    Notes:
        A 1-D Pixel
            +-----+
            ^  ^  ^
          left |  L right
               |
             center

    Example:
        >>> # Case where we have a registered func
        >>> from netharn.analytic.receptive_field_for import *
        >>> self = nn.Sequential(
        >>>     nn.Conv2d(2, 3, kernel_size=3),
        >>>     nn.Conv2d(3, 5, kernel_size=3),
        >>> )
        >>> rfield = ReceptiveFieldFor(self)()
        >>> print('rfield.hidden = {}'.format(ub.repr2(rfield.hidden, nl=3, with_dtype=False)))
        >>> print('rfield = {}'.format(ub.repr2(rfield, nl=1, with_dtype=False)))
        rfield.hidden = {
            '0': {
                'crop': np.array([1., 1.]),
                'shape': np.array([3., 3.]),
                'stride': np.array([1., 1.]),
            },
            '1': {
                'crop': np.array([2., 2.]),
                'shape': np.array([5., 5.]),
                'stride': np.array([1., 1.]),
            },
        }
        rfield = {
            'crop': np.array([2., 2.]),
            'shape': np.array([5., 5.]),
            'stride': np.array([1., 1.]),
        }

    Example:
        >>> # Case where we haven't registered a func
        >>> self = nn.Conv2d(2, 3, kernel_size=3)
        >>> rfield = ReceptiveFieldFor(self)()
        >>> print('rfield = {}'.format(ub.repr2(rfield, nl=1, with_dtype=False)))
        rfield = {
            'crop': np.array([1., 1.]),
            'shape': np.array([3., 3.]),
            'stride': np.array([1., 1.]),
        }

    Example:
        >>> # xdoctest: +REQUIRES(--network)
        >>> import torchvision  # NOQA
        >>> module = torchvision.models.alexnet().features
        >>> field = ReceptiveFieldFor(module)()
        >>> print(ub.repr2(field, nl=1, with_dtype=False))
        {
            'crop': np.array([31., 31.]),
            'shape': np.array([195., 195.]),
            'stride': np.array([32., 32.]),
        }
    """
    # impl = math  # for hacking in sympy

    def __init__(self, module):
        self.module = module
        self._func = getattr(module, 'receptive_field_for', None)
        if self._func is None:
            # Lookup rfield func if we can't find it
            found = []
            for type, _func in REGISTERED_TYPES:
                try:
                    if module is type or isinstance(module, type):
                        found.append(_func)
                except TypeError:
                    pass
            if len(found) == 1:
                self._func = found[0]
            elif len(found) == 0:
                raise ReceptiveFieldTypeError('Unknown (rf) module type {}'.format(module))
            else:
                raise AssertionError('Ambiguous (rf) module {}. Found {}'.format(module, found))

    def __call__(self, *args, **kwargs):
        if isinstance(self.module, nn.Module):
            # bound methods dont need module
            is_bound  = hasattr(self._func, '__func__') and getattr(self._func, '__func__', None) is not None
            is_bound |= hasattr(self._func, 'im_func') and getattr(self._func, 'im_func', None) is not None
            if is_bound:
                rfield = self._func(*args, **kwargs)
            else:
                # nn.Module with state
                rfield = self._func(self.module, *args, **kwargs)
        else:
            # a simple pytorch func
            rfield = self._func(*args, **kwargs)

        rfield = ReceptiveField.coerce(rfield)
        return rfield

    # @staticmethod
    # def view(arr, *args):
    #     """
    #     Wraps view calls

    #     Example:
    #         >>> arr = (2, 32, 9, 9)
    #         >>> result = OutputShapeFor.view(arr, -1)
    #         >>> assert result == (5184,)
    #     """
    #     from netharn import layers
    #     reshape = layers.Reshape(*args)
    #     return reshape.output_shape_for(arr)

    # @staticmethod
    def shape(arr):
        """
        Wraps shape calls
        """
        raise ReceptiveFieldTypeError('RF is currently unable to inspect output shape')

    @staticmethod
    def _elementwise(field1, field2):
        # Combines two receptive fields in an elementwise fashion
        field = ReceptiveField({
            'shape': np.maximum(field1['shape'], field2['shape']),
            'crop': np.maximum(field1['crop'], field2['crop']),
            'stride': np.maximum(field1['stride'], field2['stride']),
        })
        return field

    @staticmethod
    def add(field1, field2):
        return ReceptiveFieldFor._elementwise(field1, field2)

    @staticmethod
    def mul(field1, field2):
        return ReceptiveFieldFor._elementwise(field1, field2)

    @staticmethod
    def sub(field1, field2):
        return ReceptiveFieldFor._elementwise(field1, field2)

    @staticmethod
    def div(field1, field2):
        return ReceptiveFieldFor._elementwise(field1, field2)

def effective_receptive_feild(module, inputs, output_key=None, sigma=0,
                              thresh=1.00, ignore_norms=True,
                              ignore_extra=None):
    """
    Empirically measures the effective receptive feild of a network

    Method from [0], implementation loosely based on [1].

    Args:
        module (torch.nn.Module) : the network

        inputs (torch.nn.Tensor) : the input to the network. Must share the
            same device as `module`.

        output_key (None | str | Callable): If the network outputs a non-tensor
            then this should be a function that does postprocessing and returns
            a relevant Tensor that can be used to compute gradients. If the
            output is a dictionary then this can also be a string-based key
            used to lookup the appropriate output.

        sigma (float, default=0): smoothness factor (via gaussian blur)

        thresh (float, default=1.00): only consider this fraction of the
            data as meaningful (i.e. find the effective RF shape that explains
            95% of the data). A threshold of 1.0 or greater does nothing.

        ignore_norms (bool, default=True): if True ignores normalization layers
            like batch and group norm which adds negligable, but non-zero
            impact everywhere and causes the ERF shape estimation to be
            dramatically greater than it should be (although the impact still
            makes sense).

        ignore_extra (List[type], optioanl): if specified, any layer that is a
            subclass of one of these types is also ignored.

    Returns:
        dict: containing keys
            'shape' containing the effective RF shape and
            'impact' which contains the thresholded distribution

    References:
        [0] https://arxiv.org/pdf/1701.04128.pdf
        [1] https://github.com/rogertrullo/Receptive-Field-in-Pytorch/blob/master/compute_RF.py

    Example:
        >>> from netharn.analytic.receptive_field_for import *
        >>> import torchvision  # NOQA
        >>> module = nn.Sequential(*[nn.Conv2d(1, 1, 3) for i in range(10)])
        >>> inputs = torch.rand(1, 1, 200, 200)
        >>> emperical_field = effective_receptive_feild(module, inputs)
        >>> theoretic_field = ReceptiveFieldFor(module)()
        >>> # The emperical results should never be bigger than the theoretical
        >>> assert np.all(emperical_field['shape'] <= theoretic_field['shape'])

        >>> # xdoctest: +REQUIRES(--slow)
        >>> module = torchvision.models.alexnet().features
        >>> inputs = torch.rand(1, 3, 224, 224)
        >>> emperical_field = effective_receptive_feild(module, inputs)
        >>> theoretic_field = ReceptiveFieldFor(module)()
        >>> # The emperical results should never be bigger than the theoretical
        >>> assert np.all(emperical_field['shape'] <= theoretic_field['shape'])

        >>> # xdoctest: +REQUIRES(--slow)
        >>> import netharn as nh
        >>> xpu = nh.XPU.coerce('auto')
        >>> module = xpu.move(torchvision.models.vgg11_bn().features)
        >>> inputs = xpu.move(torch.rand(1, 3, 224, 224))
        >>> emperical_field = effective_receptive_feild(module, inputs)
        >>> theoretic_field = ReceptiveFieldFor(module)()
        >>> # The emperical results should never be bigger than the theoretical
        >>> assert np.all(emperical_field['shape'] <= theoretic_field['shape'])

        >>> # xdoctest: +REQUIRES(--show)
        >>> import kwplot
        >>> kwplot.autompl()
        >>> kwplot.imshow(emperical_field['impact'], doclf=True)

    Ignore:
        >>> def forward(self, x):
        >>>     # See note [TorchScript super()]
        >>>     x = self.conv1(x)
        >>>     x = self.bn1(x)
        >>>     x = self.relu(x)
        >>>     x = self.maxpool(x)
        >>> #
        >>>     x = self.layer1(x)
        >>>     x = self.layer2(x)
        >>>     x = self.layer3(x)
        >>>     x = self.layer4(x)
        >>> #
        >>>     #x = self.avgpool(x)
        >>>     #x = torch.flatten(x, 1)
        >>>     #x = self.fc(x)
        >>>     return x
        >>> xpu = nh.XPU.coerce('auto')
        >>> module1 = torchvision.models.resnet50()
        >>> ub.inject_method(module1, forward)
        >>> module1 = xpu.move(module1)
        >>> module2 = torchvision.models.resnet50(pretrained=True)
        >>> module2 = xpu.move(module2)
        >>> ub.inject_method(module2, forward)
        >>> import kwimage
        >>> img = kwimage.grab_test_image(dsize=(224, 224))
        >>> inputs = torch.from_numpy(img.transpose(2, 0, 1)[None, :] / 255.).float()
        >>> inputs = xpu.move(inputs)
        >>> #inputs = xpu.move(torch.rand(8, 3, 224, 224))
        >>> ignore_norms = 1
        >>> emperical_field1 = effective_receptive_feild(module1, inputs, ignore_norms=ignore_norms)
        >>> emperical_field2 = effective_receptive_feild(module2, inputs, ignore_norms=ignore_norms)
        >>> import kwplot
        >>> kwplot.autompl()
        >>> kwplot.imshow(inputs[0].data.cpu().numpy(), fnum=1, pnum=(1, 3, 1), title='input', doclf=1)
        >>> kwplot.imshow(emperical_field1['impact'], fnum=1, pnum=(1, 3, 2), title='pretrained=False')
        >>> kwplot.imshow(emperical_field2['impact'], doclf=0, fnum=1, pnum=(1, 3, 3), title='pretrained=True')
    """
    import netharn as nh

    # zero gradients
    for p in module.parameters():
        if p.grad is not None:
            p.grad.detach_()
            p.grad.zero_()

    if inputs.grad is not None:
        inputs.grad.detach_()
        inputs.grad.zero_()

    inputs.requires_grad = True
    # if inputs.grad is not None:
    #     raise ValueError('inputs alread has accumulated gradients')

    # Completely ignore BatchNorm layers as they will give the entire input
    # some negligable but non-zero effect on the receptive feild.
    ignored = []
    if ignore_norms:
        ignored += [
            nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d,
            nn.modules.normalization.GroupNorm,
            nn.modules.normalization.LocalResponseNorm,
            nn.modules.normalization.LayerNorm, nn.CrossMapLRN2d,
            nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d,
            nh.layers.L2Norm,
        ]
    if ignore_extra:
        ignored += ignore_extra
    with nh.util.IgnoreLayerContext(module, tuple(ignored)):
        outputs = module(inputs)

    # Note: grab a single (likely FCN) output channel
    if output_key is None:
        output_y = outputs
    elif callable(output_key):
        output_y = output_key(outputs)
    else:
        output_y = outputs[output_key]
    # elif isinstance(output_key, (six.string_types, int)):
    # else:
    #     raise TypeError('output_key={} is not understood'.format(output_key))

    if not isinstance(output_y, torch.Tensor):
        raise TypeError(
            'The output is a {}, not a tensor. Please specify '
            'output_key and ensure it returns a Tensor.'.format(type(outputs)))

    # Note: this still does the right thing if there is no spatial component.
    # because all outputs are center outputs.
    center_dims = (np.array(output_y.shape[2:]) // 2).tolist()
    center_slice = [slice(None), slice(None)] + center_dims

    # We dont need to compute a loss because we can explicitly set gradients.
    # Yay torch!
    # Explicilty set ∂l/∂y[:] = 0
    # Explicilty set ∂l/∂y[center] = 1
    grad_loss_wrt_y = torch.zeros_like(output_y)
    grad_loss_wrt_y[...] = 0
    grad_loss_wrt_y[center_slice] = 1

    # Backpropogate as if the grad of the loss wrt to y[center] was 1.
    # Note: this can take a long time on the CPU (sometimes?)
    output_y.backward(gradient=grad_loss_wrt_y)

    # The input gradient is now a measure of how much it can impact the output.
    impact = inputs.grad.abs()

    # Average the impact over all batches and all channels
    average_impact = impact.mean(dim=0).mean(dim=0)

    if isinstance(average_impact, torch.Tensor):
        average_impact = average_impact.data.cpu().numpy()

    idx_nonzeros = np.where(average_impact != 0)
    rf_bounds = [(0, 0) if len(idx) == 0 else (idx.min(), idx.max()) for idx in idx_nonzeros]
    rf_shape = [(mx - mn + 1) for mn, mx in rf_bounds]
    rf_slice = tuple([slice(mn, mx + 1) for mn, mx in rf_bounds])

    # Crop out the average impact zone for visualization
    # Normalize to have a maximum value of 1.0
    rf_impact = average_impact[rf_slice]
    rf_impact /= rf_impact.max()

    rf_impact = torch.FloatTensor(rf_impact)
    if sigma > 0:
        # Smooth things out
        _blur = nh.layers.GaussianBlurNd(dim=1, num_features=1, sigma=sigma)
        _blur.to(rf_impact.device)
        rf_impact = _blur(rf_impact[None, None])[0, 0]

    if thresh < 1:
        density = rf_impact.contiguous().view(-1).cpu().numpy().copy()
        density.sort()
        density = density[::-1]
        # Find the value threshold that explains thresh (e.g. 95%) of the data
        idx = np.where(density.cumsum() > thresh * density.sum())[0]
        lowval = float(density[idx[0]])

        effective_impact = rf_impact * (rf_impact > lowval).float()
        effective_idx_nonzeros = np.where(effective_impact != 0)
        effective_rf_bounds = [(idx.min(), idx.max()) for idx in effective_idx_nonzeros]
        effective_shape = [(mx - mn + 1) for mn, mx in effective_rf_bounds]
    else:
        effective_impact = rf_impact
        effective_rf_bounds = rf_shape
        effective_shape = rf_shape

    emperical_field = {
        'shape': effective_shape,
        'impact': effective_impact,
        'thresh': thresh,
    }
    return emperical_field

if __name__ == '__main__':
    """
    CommandLine:
        xdoctest -m netharn.analytic.receptive_field_for all --network
    """
    import xdoctest
    xdoctest.doctest_module(__file__)


#!/usr/bin/env python
# CREATED:2014-01-18 14:09:05 by <NAME> <<EMAIL>>
# unit tests for util routines

# Disable cache
import os

try:
    os.environ.pop("LIBROSA_CACHE_DIR")
except:
    pass

import platform
import numpy as np
import scipy.sparse
import pytest
import warnings
import librosa

from test_core import srand

np.set_printoptions(precision=3)

def test_example_audio_file():

    assert os.path.exists(librosa.util.example_audio_file())

@pytest.mark.parametrize("frame_length", [4, 8])
@pytest.mark.parametrize("hop_length", [2, 4])
@pytest.mark.parametrize("y", [np.random.randn(32)])
@pytest.mark.parametrize("axis", [0, -1])
def test_frame1d(frame_length, hop_length, axis, y):

    y_frame = librosa.util.frame(y, frame_length=frame_length, hop_length=hop_length, axis=axis)

    if axis == -1:
        y_frame = y_frame.T

    for i in range(y_frame.shape[0]):
        assert np.allclose(y_frame[i], y[i * hop_length : (i * hop_length + frame_length)])

@pytest.mark.parametrize("frame_length", [4, 8])
@pytest.mark.parametrize("hop_length", [2, 4])
@pytest.mark.parametrize(
    "y, axis", [(np.asfortranarray(np.random.randn(16, 32)), -1), (np.ascontiguousarray(np.random.randn(16, 32)), 0)]
)
def test_frame2d(frame_length, hop_length, axis, y):

    y_frame = librosa.util.frame(y, frame_length=frame_length, hop_length=hop_length, axis=axis)

    if axis == -1:
        y_frame = y_frame.T
        y = y.T

    for i in range(y_frame.shape[0]):
        assert np.allclose(y_frame[i], y[i * hop_length : (i * hop_length + frame_length)])

def test_frame_0stride():
    x = np.arange(10)
    xpad = x[np.newaxis]

    xpad2 = np.atleast_2d(x)

    xf = librosa.util.frame(x, 3, 1)
    xfpad = librosa.util.frame(xpad, 3, 1)
    xfpad2 = librosa.util.frame(xpad2, 3, 1)

    assert np.allclose(xf, xfpad)
    assert np.allclose(xf, xfpad2)

@pytest.mark.xfail(raises=librosa.ParameterError)
def test_frame_badtype():
    librosa.util.frame([1, 2, 3, 4], frame_length=2, hop_length=1)

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize("axis", [0, -1])
@pytest.mark.parametrize("x", [np.arange(16)])
def test_frame_too_short(x, axis):
    librosa.util.frame(x, frame_length=17, hop_length=1, axis=axis)

@pytest.mark.xfail(raises=librosa.ParameterError)
def test_frame_bad_hop():
    librosa.util.frame(np.arange(16), frame_length=4, hop_length=0)

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize("axis", [1, 2])
def test_frame_bad_axis(axis):
    librosa.util.frame(np.zeros((3, 3, 3)), frame_length=2, hop_length=1, axis=axis)

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize("x, axis", [(np.zeros((4, 4), order="C"), -1), (np.zeros((4, 4), order="F"), 0)])
def test_frame_bad_contiguity(x, axis):
    librosa.util.frame(x, frame_length=2, hop_length=1, axis=axis)

@pytest.mark.parametrize("y", [np.ones((16,)), np.ones((16, 16))])
@pytest.mark.parametrize("m", [0, 10])
@pytest.mark.parametrize("axis", [0, -1])
@pytest.mark.parametrize("mode", ["constant", "edge", "reflect"])
def test_pad_center(y, m, axis, mode):

    n = m + y.shape[axis]
    y_out = librosa.util.pad_center(y, n, axis=axis, mode=mode)

    n_len = y.shape[axis]
    n_pad = int((n - n_len) / 2)

    eq_slice = [slice(None)] * y.ndim
    eq_slice[axis] = slice(n_pad, n_pad + n_len)

    assert np.allclose(y, y_out[tuple(eq_slice)])

@pytest.mark.parametrize("y", [np.ones((16,)), np.ones((16, 16))])
@pytest.mark.parametrize("n", [0, 10])
@pytest.mark.parametrize("axis", [0, -1])
@pytest.mark.parametrize("mode", ["constant", "edge", "reflect"])
@pytest.mark.xfail(raises=librosa.ParameterError)
def test_pad_center_fail(y, n, axis, mode):
    librosa.util.pad_center(y, n, axis=axis, mode=mode)

@pytest.mark.parametrize("y", [np.ones((16,)), np.ones((16, 16))])
@pytest.mark.parametrize("m", [-5, 0, 5])
@pytest.mark.parametrize("axis", [0, -1])
def test_fix_length(y, m, axis):
    n = m + y.shape[axis]

    y_out = librosa.util.fix_length(y, n, axis=axis)

    eq_slice = [slice(None)] * y.ndim
    eq_slice[axis] = slice(y.shape[axis])

    if n > y.shape[axis]:
        assert np.allclose(y, y_out[tuple(eq_slice)])
    else:
        assert np.allclose(y[tuple(eq_slice)], y)

@pytest.mark.parametrize("frames", [np.arange(20, 100, step=15)])
@pytest.mark.parametrize("x_min", [0, 20])
@pytest.mark.parametrize("x_max", [20, 70, 120])
@pytest.mark.parametrize("pad", [False, True])
def test_fix_frames(frames, x_min, x_max, pad):

    f_fix = librosa.util.fix_frames(frames, x_min=x_min, x_max=x_max, pad=pad)

    if x_min is not None:
        if pad:
            assert f_fix[0] == x_min
        assert np.all(f_fix >= x_min)

    if x_max is not None:
        if pad:
            assert f_fix[-1] == x_max
        assert np.all(f_fix <= x_max)

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize("frames", [np.arange(-20, 100)])
@pytest.mark.parametrize("x_min", [None, 0, 20])
@pytest.mark.parametrize("x_max", [None, 0, 20])
@pytest.mark.parametrize("pad", [False, True])
def test_fix_frames_fail_negative(frames, x_min, x_max, pad):
    librosa.util.fix_frames(frames, x_min, x_max, pad)

@pytest.mark.parametrize("norm", [np.inf, -np.inf, 0, 0.5, 1.0, 2.0, None])
@pytest.mark.parametrize("ndims,axis", [(1, 0), (1, -1), (2, 0), (2, 1), (2, -1), (3, 0), (3, 1), (3, 2), (3, -1)])
def test_normalize(ndims, norm, axis):
    srand()
    X = np.random.randn(*([4] * ndims))
    X_norm = librosa.util.normalize(X, norm=norm, axis=axis)

    # Shape and dtype checks
    assert X_norm.dtype == X.dtype
    assert X_norm.shape == X.shape

    if norm is None:
        assert np.allclose(X, X_norm)
        return

    X_norm = np.abs(X_norm)

    if norm == np.inf:
        values = np.max(X_norm, axis=axis)
    elif norm == -np.inf:
        values = np.min(X_norm, axis=axis)
    elif norm == 0:
        # XXX: normalization here isn't quite right
        values = np.ones(1)

    else:
        values = np.sum(X_norm ** norm, axis=axis) ** (1.0 / norm)

    assert np.allclose(values, np.ones_like(values))

@pytest.mark.parametrize("norm", ["inf", -0.5, -2])
@pytest.mark.parametrize("X", [np.ones((3, 3))])
@pytest.mark.xfail(raises=librosa.ParameterError)
def test_normalize_badnorm(X, norm):
    librosa.util.normalize(X, norm=norm)

@pytest.mark.parametrize("badval", [np.nan, np.inf, -np.inf])
@pytest.mark.xfail(raises=librosa.ParameterError)
def test_normalize_bad_input(badval):
    X = np.ones((3, 3))
    X[0] = badval
    librosa.util.normalize(X, norm=np.inf, axis=0)

@pytest.mark.parametrize("fill", [7, "foo"])
@pytest.mark.parametrize("X", [np.ones((2, 2))])
@pytest.mark.xfail(raises=librosa.ParameterError)
def test_normalize_badfill(X, fill):
    librosa.util.normalize(X, fill=fill)

@pytest.mark.parametrize("x", [np.asarray([[0, 1, 2, 3]])])
@pytest.mark.parametrize(
    "threshold, result",
    [(None, [[0, 1, 1, 1]]), (1, [[0, 1, 1, 1]]), (2, [[0, 1, 1, 1]]), (3, [[0, 1, 2, 1]]), (4, [[0, 1, 2, 3]])],
)
def test_normalize_threshold(x, threshold, result):
    assert np.allclose(librosa.util.normalize(x, threshold=threshold), result)

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize("x", [np.asarray([[0, 1, 2, 3]])])
@pytest.mark.parametrize("threshold", [0, -1])
def test_normalize_threshold_fail(x, threshold):
    librosa.util.normalize(x, threshold=threshold)

@pytest.mark.xfail(raises=librosa.ParameterError)
def test_normalize_fill_l0():
    X = np.ones((2, 2))
    librosa.util.normalize(X, fill=True, norm=0)

@pytest.mark.parametrize("norm", [1, 2, np.inf])
@pytest.mark.parametrize("X", [np.zeros((3, 3))])
def test_normalize_fill_allaxes(X, norm):
    Xn = librosa.util.normalize(X, fill=True, axis=None, norm=norm)
    if norm is np.inf:
        assert np.allclose(Xn, 1)
    else:
        assert np.allclose(np.sum(Xn ** norm) ** (1.0 / norm), 1)

@pytest.mark.parametrize("norm", [1, 2, np.inf])
@pytest.mark.parametrize("X", [np.zeros((3, 3))])
def test_normalize_nofill(X, norm):
    Xn = librosa.util.normalize(X, fill=False, norm=norm)
    assert np.allclose(Xn, 0)

@pytest.mark.parametrize("X", [np.asarray([[0.0, 1], [0, 1]])])
@pytest.mark.parametrize("norm,value", [(1, 0.5), (2, np.sqrt(2) / 2), (np.inf, 1)])
@pytest.mark.parametrize("threshold", [0.5, 2])
def test_normalize_fill(X, threshold, norm, value):
    Xn = librosa.util.normalize(X, fill=True, norm=norm, threshold=threshold)
    assert np.allclose(Xn, value)

@pytest.mark.parametrize("ndim", [1, 3])
@pytest.mark.parametrize("axis", [0, 1, -1])
@pytest.mark.parametrize("index", [False, True])
@pytest.mark.parametrize("value", [None, np.min, np.mean, np.max])
@pytest.mark.xfail(raises=librosa.ParameterError)
def test_axis_sort_badndim(ndim, axis, index, value):
    data = np.zeros([2] * ndim)
    librosa.util.axis_sort(data, axis=axis, index=index, value=value)

@pytest.mark.parametrize("ndim", [2])
@pytest.mark.parametrize("axis", [0, 1, -1])
@pytest.mark.parametrize("index", [False, True])
@pytest.mark.parametrize("value", [None, np.min, np.mean, np.max])
def test_axis_sort(ndim, axis, index, value):
    srand()
    data = np.random.randn(*([10] * ndim))
    if index:
        Xsorted, idx = librosa.util.axis_sort(data, axis=axis, index=index, value=value)

        cmp_slice = [slice(None)] * ndim
        cmp_slice[axis] = idx

        assert np.allclose(data[tuple(cmp_slice)], Xsorted)

    else:
        Xsorted = librosa.util.axis_sort(data, axis=axis, index=index, value=value)

    compare_axis = np.mod(1 - axis, 2)

    if value is None:
        value = np.argmax

    sort_values = value(Xsorted, axis=compare_axis)

    assert np.allclose(sort_values, np.sort(sort_values))

@pytest.mark.parametrize(
    "int_from, int_to",
    [
        (np.asarray([[0, 2], [0, 4], [3, 6]]), np.empty((0, 2), dtype=int)),
        (np.empty((0, 2), dtype=int), np.asarray([[0, 2], [0, 4], [3, 6]])),
    ],
)
@pytest.mark.xfail(raises=librosa.ParameterError)
def test_match_intervals_empty(int_from, int_to):
    librosa.util.match_intervals(int_from, int_to)

@pytest.mark.xfail(raises=librosa.ParameterError)
def test_match_intervals_strict_fail():
    int_from = np.asarray([[0, 3], [2, 4], [5, 7]])
    int_to = np.asarray([[0, 2], [0, 4]])
    librosa.util.match_intervals(int_from, int_to, strict=True)

@pytest.mark.parametrize("int_from", [np.asarray([[0, 3], [2, 4], [5, 7]])])
@pytest.mark.parametrize("int_to", [np.asarray([[0, 2], [0, 4], [3, 6]])])
@pytest.mark.parametrize("matches", [np.asarray([1, 1, 2])])
def test_match_intervals_strict(int_from, int_to, matches):

    test_matches = librosa.util.match_intervals(int_from, int_to, strict=True)
    assert np.array_equal(matches, test_matches)

@pytest.mark.parametrize("int_from", [np.asarray([[0, 3], [2, 4], [5, 7]])])
@pytest.mark.parametrize(
    "int_to,matches",
    [
        (np.asarray([[0, 2], [0, 4], [3, 6]]), np.asarray([1, 1, 2])),
        (np.asarray([[0, 2], [0, 4]]), np.asarray([1, 1, 1])),
    ],
)
def test_match_intervals_nonstrict(int_from, int_to, matches):
    test_matches = librosa.util.match_intervals(int_from, int_to, strict=False)
    assert np.array_equal(matches, test_matches)

@pytest.mark.parametrize("n", [1, 5, 20, 100])
@pytest.mark.parametrize("m", [1, 5, 20, 100])
def test_match_events(n, m):

    srand()
    ev1 = np.abs(np.random.randn(n))
    ev2 = np.abs(np.random.randn(m))

    match = librosa.util.match_events(ev1, ev2)

    for i in range(len(match)):
        values = np.asarray([np.abs(ev1[i] - e2) for e2 in ev2])
        assert not np.any(values < values[match[i]])

@pytest.mark.parametrize("ev1,ev2", [(np.array([]), np.arange(5)), (np.arange(5), np.array([]))])
@pytest.mark.xfail(raises=librosa.ParameterError)
def test_match_events_failempty(ev1, ev2):
    librosa.util.match_events(ev1, ev2)

@pytest.mark.parametrize("events_from", [np.asarray([5, 15, 25])])
@pytest.mark.parametrize("events_to", [np.asarray([0, 10, 20, 30])])
@pytest.mark.parametrize("left,right,target", [(False, True, [10, 20, 30]), (True, False, [0, 10, 20])])
def test_match_events_onesided(events_from, events_to, left, right, target):

    events_from = np.asarray(events_from)
    events_to = np.asarray(events_to)
    match = librosa.util.match_events(events_from, events_to, left=left, right=right)

    assert np.allclose(target, events_to[match])

def test_match_events_twosided():
    events_from = np.asarray([5, 15, 25])
    events_to = np.asarray([5, 15, 25, 30])
    match = librosa.util.match_events(events_from, events_to, left=False, right=False)
    assert np.allclose(match, [0, 1, 2])

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize(
    "events_from,events_to,left,right",
    [
        ([40, 15, 25], [0, 10, 20, 30], False, True),  # right-sided fail
        ([-1, 15, 25], [0, 10, 20, 30], True, False),  # left-sided fail
        ([-1, 15, 25], [0, 10, 20, 30], False, False),  # two-sided fail
    ],
)
def test_match_events_onesided_fail(events_from, events_to, left, right):
    events_from = np.asarray(events_from)
    events_to = np.asarray(events_to)
    librosa.util.match_events(events_from, events_to, left=left, right=right)

@pytest.mark.parametrize("ndim, axis", [(n, m) for n in range(1, 5) for m in range(n)])
def test_localmax(ndim, axis):

    srand()

    data = np.random.randn(*([7] * ndim))
    lm = librosa.util.localmax(data, axis=axis)

    for hits in np.argwhere(lm):
        for offset in [-1, 1]:
            compare_idx = hits.copy()
            compare_idx[axis] += offset

            if compare_idx[axis] < 0:
                continue

            if compare_idx[axis] >= data.shape[axis]:
                continue

            if offset < 0:
                assert data[tuple(hits)] > data[tuple(compare_idx)]
            else:
                assert data[tuple(hits)] >= data[tuple(compare_idx)]

@pytest.mark.parametrize("x", [np.random.randn(_) ** 2 for _ in [1, 5, 10, 100]])
@pytest.mark.parametrize("pre_max", [0, 1, 10])
@pytest.mark.parametrize("post_max", [1, 10])
@pytest.mark.parametrize("pre_avg", [0, 1, 10])
@pytest.mark.parametrize("post_avg", [1, 10])
@pytest.mark.parametrize("wait", [0, 1, 10])
@pytest.mark.parametrize("delta", [0.05, 100.0])
def test_peak_pick(x, pre_max, post_max, pre_avg, post_avg, delta, wait):
    peaks = librosa.util.peak_pick(x, pre_max, post_max, pre_avg, post_avg, delta, wait)

    for i in peaks:
        # Test 1: is it a peak in this window?
        s = i - pre_max
        if s < 0:
            s = 0
        t = i + post_max

        diff = x[i] - np.max(x[s:t])
        assert diff > 0 or np.isclose(diff, 0, rtol=1e-3, atol=1e-4)

        # Test 2: is it a big enough peak to count?
        s = i - pre_avg
        if s < 0:
            s = 0
        t = i + post_avg

        diff = x[i] - (delta + np.mean(x[s:t]))
        assert diff > 0 or np.isclose(diff, 0, rtol=1e-3, atol=1e-4)

    # Test 3: peak separation
    assert not np.any(np.diff(peaks) <= wait)

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize("x", [np.random.randn(_) ** 2 for _ in [1, 5, 10, 100]])
@pytest.mark.parametrize(
    "pre_max,post_max,pre_avg,post_avg,delta,wait",
    [
        (-1, 1, 1, 1, 0.05, 1),  # negative pre-max
        (1, -1, 1, 1, 0.05, 1),  # negative post-max
        (1, 0, 1, 1, 0.05, 1),  # 0 post-max
        (1, 1, -1, 1, 0.05, 1),  # negative pre-avg
        (1, 1, 1, -1, 0.05, 1),  # negative post-avg
        (1, 1, 1, 0, 0.05, 1),  # zero post-avg
        (1, 1, 1, 1, -0.05, 1),  # negative delta
        (1, 1, 1, 1, 0.05, -1),  # negative wait
    ],
)
def test_peak_pick_fail(x, pre_max, post_max, pre_avg, post_avg, delta, wait):
    librosa.util.peak_pick(x, pre_max, post_max, pre_avg, post_avg, delta, wait)

@pytest.mark.xfail(raises=librosa.ParameterError)
def test_peak_pick_shape_fail():
    # Can't pick peaks on 2d inputs
    librosa.util.peak_pick(np.eye(2), 1, 1, 1, 1, 0.5, 1)

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize("ndim", [3, 4])
def test_sparsify_rows_ndimfail(ndim):
    X = np.zeros([2] * ndim)
    librosa.util.sparsify_rows(X)

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize("quantile", [1.0, -1, 2.0])
@pytest.mark.parametrize("X", [np.ones((3, 3))])
def test_sparsify_rows_badquantile(X, quantile):
    librosa.util.sparsify_rows(X, quantile=quantile)

@pytest.mark.parametrize("ndim", [1, 2])
@pytest.mark.parametrize("d", [1, 5, 10, 100])
@pytest.mark.parametrize("q", [0.0, 0.01, 0.25, 0.5, 0.99])
def test_sparsify_rows(ndim, d, q):
    srand()

    X = np.random.randn(*([d] * ndim)) ** 4

    X = np.asarray(X)

    xs = librosa.util.sparsify_rows(X, quantile=q)

    if ndim == 1:
        X = X.reshape((1, -1))

    assert np.allclose(xs.shape, X.shape)

    # And make sure that xs matches X on nonzeros
    xsd = np.asarray(xs.todense())

    for i in range(xs.shape[0]):
        assert np.allclose(xsd[i, xs[i].indices], X[i, xs[i].indices])

    # Compute row-wise magnitude marginals
    v_in = np.sum(np.abs(X), axis=-1)
    v_out = np.sum(np.abs(xsd), axis=-1)

    # Ensure that v_out retains 1-q fraction of v_in
    assert np.all(v_out >= (1.0 - q) * v_in)

@pytest.mark.parametrize(
    "searchdir", [os.path.join(os.path.curdir, "tests"), os.path.join(os.path.curdir, "tests", "data")]
)
@pytest.mark.parametrize("ext", [None, "wav", "WAV", ["wav"], ["WAV"]])
@pytest.mark.parametrize("recurse", [True])
@pytest.mark.parametrize("case_sensitive", list({False} | {platform.system() != "Windows"}))
@pytest.mark.parametrize("limit", [None, 1, 2])
@pytest.mark.parametrize("offset", [0, 1, -1])
@pytest.mark.parametrize(
    "output",
    [
        [
            os.path.join(os.path.abspath(os.path.curdir), "tests", "data", s)
            for s in ["test1_22050.mp3", "test1_22050.wav", "test1_44100.wav", "test2_8000.wav"]
        ]
    ],
)
def test_find_files(searchdir, ext, recurse, case_sensitive, limit, offset, output):
    files = librosa.util.find_files(
        searchdir, ext=ext, recurse=recurse, case_sensitive=case_sensitive, limit=limit, offset=offset
    )

    targets = output
    if ext is not None:
        # If we're only seeking wavs, bump off the mp3 file
        targets = targets[1:]

    s1 = slice(offset, None)
    s2 = slice(limit)

    if case_sensitive and ext not in (None, "wav", ["wav"]):
        assert len(files) == 0
    else:
        assert set(files) == set(targets[s1][s2])

def test_find_files_nonrecurse():
    files = librosa.util.find_files(os.path.join(os.path.curdir, "tests"), recurse=False)
    assert len(files) == 0

# fail if ext is not none, we're case-sensitive, and looking for WAV
@pytest.mark.parametrize("ext", ["WAV", ["WAV"]])
def test_find_files_case_sensitive(ext):
    files = librosa.util.find_files(os.path.join(os.path.curdir, "tests"), ext=ext, case_sensitive=True)
    # On windows, this test won't work
    if platform.system() != "Windows":
        assert len(files) == 0

@pytest.mark.parametrize("x_in", np.linspace(-2, 2, num=6))
@pytest.mark.parametrize("cast", [None, np.floor, np.ceil])
def test_valid_int(x_in, cast):

    z = librosa.util.valid_int(x_in, cast)

    assert isinstance(z, int)
    if cast is None:
        assert z == int(np.floor(x_in))
    else:
        assert z == int(cast(x_in))

@pytest.mark.parametrize("x", np.linspace(-2, 2, num=3))
@pytest.mark.parametrize("cast", [7])
@pytest.mark.xfail(raises=librosa.ParameterError)
def test_valid_int_fail(x, cast):
    # Test with a non-callable cast operator
    librosa.util.valid_int(x, cast)

@pytest.mark.parametrize(
    "ivals", [np.asarray([[0, 1], [1, 2]]), np.asarray([[0, 0], [1, 1]]), np.asarray([[0, 2], [1, 2]])]
)
def test_valid_intervals(ivals):
    librosa.util.valid_intervals(ivals)

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize(
    "ivals", [np.asarray([]), np.arange(2), np.ones((2, 2, 2)), np.ones((2, 3))]  # ndim=0  # ndim=1  # ndim=3
)  # ndim=2, shape[1] != 2
def test_valid_intervals_badshape(ivals):
    #   fail if ndim != 2 or shape[1] != 2
    librosa.util.valid_intervals(ivals)

@pytest.mark.xfail(raises=librosa.ParameterError)
@pytest.mark.parametrize("intval", [np.asarray([[0, 1], [2, 1]])])
def test_valid_intervals_fail(intval):
    # Test for issue #712: intervals must have non-negative duration
    librosa.util.valid_intervals(intval)

def test_warning_deprecated():
    @librosa.util.decorators.deprecated("old_version", "new_version")
    def __dummy():
        return True

    with warnings.catch_warnings(record=True) as out:
        x = __dummy()

        # Make sure we still get the right value
        assert x is True

        # And that the warning triggered
        assert len(out) > 0

        # And that the category is correct
        assert out[0].category is DeprecationWarning

        # And that it says the right thing (roughly)
        assert "deprecated" in str(out[0].message).lower()

def test_warning_moved():
    @librosa.util.decorators.moved("from", "old_version", "new_version")
    def __dummy():
        return True

    with warnings.catch_warnings(record=True) as out:
        x = __dummy()

        # Make sure we still get the right value
        assert x is True

        # And that the warning triggered
        assert len(out) > 0

        # And that the category is correct
        assert out[0].category is DeprecationWarning

        # And that it says the right thing (roughly)
        assert "moved" in str(out[0].message).lower()

def test_warning_rename_kw_pass():

    warnings.resetwarnings()
    warnings.simplefilter("always")

    ov = librosa.util.Deprecated()
    nv = 23

    with warnings.catch_warnings(record=True) as out:
        v = librosa.util.rename_kw("old", ov, "new", nv, "0", "1")

        assert v == nv

        # Make sure no warning triggered
        assert len(out) == 0

def test_warning_rename_kw_fail():

    warnings.resetwarnings()
    warnings.simplefilter("always")

    ov = 27
    nv = 23

    with warnings.catch_warnings(record=True) as out:
        v = librosa.util.rename_kw("old", ov, "new", nv, "0", "1")

        assert v == ov

        # Make sure the warning triggered
        assert len(out) > 0

        # And that the category is correct
        assert out[0].category is DeprecationWarning

        # And that it says the right thing (roughly)
        assert "renamed" in str(out[0].message).lower()

@pytest.mark.parametrize("idx", [np.arange(10, 90, 10), np.arange(10, 90, 15)])
@pytest.mark.parametrize("idx_min", [None, 5, 15])
@pytest.mark.parametrize("idx_max", [None, 85, 100])
@pytest.mark.parametrize("step", [None, 2])
@pytest.mark.parametrize("pad", [False, True])
def test_index_to_slice(idx, idx_min, idx_max, step, pad):

    slices = librosa.util.index_to_slice(idx, idx_min=idx_min, idx_max=idx_max, step=step, pad=pad)

    if pad:
        if idx_min is not None:
            assert slices[0].start == idx_min
            if idx.min() != idx_min:
                slices = slices[1:]
        if idx_max is not None:
            assert slices[-1].stop == idx_max
            if idx.max() != idx_max:
                slices = slices[:-1]

    if idx_min is not None:
        idx = idx[idx >= idx_min]

    if idx_max is not None:
        idx = idx[idx <= idx_max]

    idx = np.unique(idx)
    assert len(slices) == len(idx) - 1

    for sl, start, stop in zip(slices, idx, idx[1:]):
        assert sl.start == start
        assert sl.stop == stop
        assert sl.step == step

@pytest.mark.parametrize("aggregate", [None, np.mean, np.sum])
@pytest.mark.parametrize("ndim,axis", [(1, 0), (1, -1), (2, 0), (2, 1), (2, -1), (3, 0), (3, 2), (3, -1)])
def test_sync(aggregate, ndim, axis):
    data = np.ones([6] * ndim, dtype=np.float)

    # Make some slices that don't fill the entire dimension
    slices = [slice(1, 3), slice(3, 4)]
    dsync = librosa.util.sync(data, slices, aggregate=aggregate, axis=axis)

    # Check the axis shapes
    assert dsync.shape[axis] == len(slices)

    s_test = list(dsync.shape)
    del s_test[axis]
    s_orig = list(data.shape)
    del s_orig[axis]
    assert s_test == s_orig

    # The first slice will sum to 2 and have mean 1
    idx = [slice(None)] * ndim
    idx[axis] = 0
    if aggregate is np.sum:
        assert np.allclose(dsync[idx], 2)
    else:
        assert np.allclose(dsync[idx], 1)

    # The second slice will sum to 1 and have mean 1
    idx[axis] = 1
    assert np.allclose(dsync[idx], 1)

@pytest.mark.parametrize("aggregate", [np.mean, np.max])
def test_sync_slices(aggregate):
    x = np.arange(8, dtype=float)
    slices = [slice(0, 2), slice(2, 4), slice(4, 6), slice(6, 8)]
    xsync = librosa.util.sync(x, slices, aggregate=aggregate)
    if aggregate is np.mean:
        assert np.allclose(xsync, [0.5, 2.5, 4.5, 6.5])
    elif aggregate is np.max:
        assert np.allclose(xsync, [1, 3, 5, 7])
    else:
        assert False

@pytest.mark.parametrize("aggregate", [np.mean, np.max])
@pytest.mark.parametrize("atype", [list, np.asarray])
def test_sync_frames(aggregate, atype):
    x = np.arange(8, dtype=float)
    frames = atype([0, 2, 4, 6, 8])
    xsync = librosa.util.sync(x, frames, aggregate=aggregate)
    if aggregate is np.mean:
        assert np.allclose(xsync, [0.5, 2.5, 4.5, 6.5])
    elif aggregate is np.max:
        assert np.allclose(xsync, [1, 3, 5, 7])
    else:
        assert False

@pytest.mark.parametrize("atype", [list, np.asarray])
@pytest.mark.parametrize("pad", [False, True])
def test_sync_frames_pad(atype, pad):
    x = np.arange(8, dtype=float)
    frames = atype([2, 4, 6])
    xsync = librosa.util.sync(x, frames, pad=pad)
    if pad:
        assert np.allclose(xsync, [0.5, 2.5, 4.5, 6.5])
    else:
        assert np.allclose(xsync, [2.5, 4.5])

@pytest.mark.parametrize("data", [np.mod(np.arange(135), 5)])
@pytest.mark.parametrize("idx", [["foo", "bar"], [None], [slice(None), None]])
@pytest.mark.xfail(raises=librosa.ParameterError)
def test_sync_fail(data, idx):
    librosa.util.sync(data, idx)

@pytest.mark.parametrize("power", [1, 2, 50, 100, np.inf])
@pytest.mark.parametrize("split_zeros", [False, True])
def test_softmask(power, split_zeros):

    srand()

    X = np.abs(np.random.randn(10, 10))
    X_ref = np.abs(np.random.randn(10, 10))

    # Zero out some rows
    X[3, :] = 0
    X_ref[3, :] = 0

    M = librosa.util.softmask(X, X_ref, power=power, split_zeros=split_zeros)

    assert np.all(0 <= M) and np.all(M <= 1)

    if split_zeros and np.isfinite(power):
        assert np.allclose(M[3, :], 0.5)
    else:
        assert not np.any(M[3, :]), M[3]

def test_softmask_int():
    X = 2 * np.ones((3, 3), dtype=np.int32)
    X_ref = np.vander(np.arange(3))

    M1 = librosa.util.softmask(X, X_ref, power=1)
    M2 = librosa.util.softmask(X_ref, X, power=1)

    assert np.allclose(M1 + M2, 1)

@pytest.mark.parametrize(
    "x,x_ref,power,split_zeros",
    [
        (-np.ones(3), np.ones(3), 1, False),
        (np.ones(3), -np.ones(3), 1, False),
        (np.ones(3), np.ones(4), 1, False),
        (np.ones(3), np.ones(3), 0, False),
        (np.ones(3), np.ones(3), -1, False),
    ],
)
@pytest.mark.xfail(raises=librosa.ParameterError)
def test_softmask_fail(x, x_ref, power, split_zeros):
    librosa.util.softmask(x, x_ref, power=power, split_zeros=split_zeros)

@pytest.mark.parametrize(
    "x,value",
    [
        (1, np.finfo(np.float32).tiny),
        (np.ones(3, dtype=int), np.finfo(np.float32).tiny),
        (np.ones(3, dtype=np.float32), np.finfo(np.float32).tiny),
        (1.0, np.finfo(np.float64).tiny),
        (np.ones(3, dtype=np.float64), np.finfo(np.float64).tiny),
        (1j, np.finfo(np.complex128).tiny),
        (np.ones(3, dtype=np.complex64), np.finfo(np.complex64).tiny),
        (np.ones(3, dtype=np.complex128), np.finfo(np.complex128).tiny),
    ],
)
def test_tiny(x, value):
    assert value == librosa.util.tiny(x)

def test_util_fill_off_diagonal_8_8():
    # Case 1: Square matrix (N=M)
    mut_x = np.ones((8, 8))
    librosa.util.fill_off_diagonal(mut_x, 0.25)

    gt_x = np.array(
        [
            [1, 1, 0, 0, 0, 0, 0, 0],
            [1, 1, 1, 0, 0, 0, 0, 0],
            [0, 1, 1, 1, 0, 0, 0, 0],
            [0, 0, 1, 1, 1, 0, 0, 0],
            [0, 0, 0, 1, 1, 1, 0, 0],
            [0, 0, 0, 0, 1, 1, 1, 0],
            [0, 0, 0, 0, 0, 1, 1, 1],
            [0, 0, 0, 0, 0, 0, 1, 1],
        ]
    )

    assert np.array_equal(mut_x, gt_x)
    assert np.array_equal(mut_x, gt_x.T)

def test_util_fill_off_diagonal_8_12():
    # Case 2a: N!=M
    mut_x = np.ones((8, 12))
    librosa.util.fill_off_diagonal(mut_x, 0.25)

    gt_x = np.array(
        [
            [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
            [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
            [0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
            [0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
            [0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0],
            [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0],
            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
            [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1],
        ]
    )

    assert np.array_equal(mut_x, gt_x)

    # Case 2b: (N!=M).T
    mut_x = np.ones((8, 12)).T
    librosa.util.fill_off_diagonal(mut_x, 0.25)

    assert np.array_equal(mut_x, gt_x.T)

@pytest.mark.parametrize("dtype_A", [np.float32, np.float64])
@pytest.mark.parametrize("dtype_B", [np.float32, np.float64])
def test_nnls_vector(dtype_A, dtype_B):
    srand()

    # Make a random basis
    A = np.random.randn(5, 7).astype(dtype_A)

    # Make a random latent vector
    x = np.random.randn(A.shape[1]) ** 2

    B = A.dot(x).astype(dtype_B)

    x_rec = librosa.util.nnls(A, B)

    assert np.all(x_rec >= 0)
    assert np.sqrt(np.mean((B - A.dot(x_rec)) ** 2)) <= 1e-6

@pytest.mark.parametrize("dtype_A", [np.float32, np.float64])
@pytest.mark.parametrize("dtype_B", [np.float32, np.float64])
@pytest.mark.parametrize("x_size", [3, 30])
def test_nnls_matrix(dtype_A, dtype_B, x_size):
    srand()

    # Make a random basis
    A = np.random.randn(5, 7).astype(dtype_A)

    # Make a random latent matrix
    #   when x_size is 3, B is 7x3 (smaller than A)
    x = np.random.randn(A.shape[1], x_size) ** 2

    B = A.dot(x).astype(dtype_B)

    x_rec = librosa.util.nnls(A, B)

    assert np.all(x_rec >= 0)
    assert np.sqrt(np.mean((B - A.dot(x_rec)) ** 2)) <= 1e-5

@pytest.mark.parametrize("dtype_A", [np.float32, np.float64])
@pytest.mark.parametrize("dtype_B", [np.float32, np.float64])
@pytest.mark.parametrize("x_size", [16, 64, 256])
def test_nnls_multiblock(dtype_A, dtype_B, x_size):
    srand()

    # Make a random basis
    A = np.random.randn(7, 1025).astype(dtype_A)

    # Make a random latent matrix
    #   when x_size is 3, B is 7x3 (smaller than A)
    x = np.random.randn(A.shape[1], x_size) ** 2

    B = A.dot(x).astype(dtype_B)

    x_rec = librosa.util.nnls(A, B)

    assert np.all(x_rec >= 0)
    assert np.sqrt(np.mean((B - A.dot(x_rec)) ** 2)) <= 1e-4

@pytest.fixture
def psig():

    # [[0, 1, 2, 3, 4]]
    # axis=1 or -1 ==> [-1.5, 1, 1, 1, -1.5]
    # axis=0 ==> [0, 0, 0, 0, 0]
    return np.arange(0, 5, dtype=float)[np.newaxis]

@pytest.mark.parametrize("edge_order", [1, 2])
@pytest.mark.parametrize("axis", [0, 1, -1])
def test_cyclic_gradient(psig, edge_order, axis):
    grad = librosa.util.cyclic_gradient(psig, edge_order=edge_order, axis=axis)

    assert grad.shape == psig.shape
    assert grad.dtype == psig.dtype

    # Check the values
    if axis == 0:
        assert np.allclose(grad, 0)
    else:
        assert np.allclose(grad, [-1.5, 1, 1, 1, -1.5])

def test_shear_dense():

    E = np.eye(3)

    E_shear = librosa.util.shear(E, factor=1, axis=0)
    assert np.allclose(E_shear, np.asarray([[1, 0, 0], [0, 0, 1], [0, 1, 0]]))

    E_shear = librosa.util.shear(E, factor=1, axis=1)
    assert np.allclose(E_shear, np.asarray([[1, 0, 0], [0, 0, 1], [0, 1, 0]]))

    E_shear = librosa.util.shear(E, factor=-1, axis=1)
    assert np.allclose(E_shear, np.asarray([[1, 1, 1], [0, 0, 0], [0, 0, 0]]))

    E_shear = librosa.util.shear(E, factor=-1, axis=0)
    assert np.allclose(E_shear, np.asarray([[1, 0, 0], [1, 0, 0], [1, 0, 0]]))

@pytest.mark.parametrize("fmt", ["csc", "csr", "lil", "dok"])
def test_shear_sparse(fmt):
    E = scipy.sparse.identity(3, format=fmt)

    E_shear = librosa.util.shear(E, factor=1, axis=0)
    assert E_shear.format == fmt
    assert np.allclose(E_shear.toarray(), np.asarray([[1, 0, 0], [0, 0, 1], [0, 1, 0]]))

    E_shear = librosa.util.shear(E, factor=1, axis=1)
    assert E_shear.format == fmt
    assert np.allclose(E_shear.toarray(), np.asarray([[1, 0, 0], [0, 0, 1], [0, 1, 0]]))

    E_shear = librosa.util.shear(E, factor=-1, axis=1)
    assert E_shear.format == fmt
    assert np.allclose(E_shear.toarray(), np.asarray([[1, 1, 1], [0, 0, 0], [0, 0, 0]]))

    E_shear = librosa.util.shear(E, factor=-1, axis=0)
    assert E_shear.format == fmt
    assert np.allclose(E_shear.toarray(), np.asarray([[1, 0, 0], [1, 0, 0], [1, 0, 0]]))

@pytest.mark.xfail(raises=librosa.ParameterError)
def test_shear_badfactor():
    librosa.util.shear(np.eye(3), factor=None)

def test_stack_contig():
    x1 = np.ones(3)
    x2 = -np.ones(3)

    xs = librosa.util.stack([x1, x2], axis=0)

    assert xs.flags["F_CONTIGUOUS"]
    assert np.allclose(xs, [[1, 1, 1], [-1, -1, -1]])

@pytest.mark.xfail(raises=librosa.ParameterError)
def test_stack_fail_shape():
    x1 = np.ones(3)

    x2 = np.ones(2)
    librosa.util.stack([x1, x2])

@pytest.mark.xfail(raises=librosa.ParameterError)
def test_stack_fail_empty():
    librosa.util.stack([])

@pytest.mark.parametrize("axis", [0, 1, -1])
@pytest.mark.parametrize("x", [np.random.randn(5, 10, 20)])
def test_stack_consistent(x, axis):
    xs = librosa.util.stack([x, x], axis=axis)
    xsnp = np.stack([x, x], axis=axis)

    assert np.allclose(xs, xsnp)
    if axis != 0:
        assert xs.flags["C_CONTIGUOUS"]


<reponame>AVoss84/gauss_mix

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import os
from scipy.stats import wishart, multivariate_normal, norm, randint, bernoulli, beta, multinomial, gamma, dirichlet, uniform

def exp_normalize(x):
    b = x.max()
    y = np.exp(x - b)
    return y / y.sum()

class mvt2mixture:

    def __init__(self, thetas = {'mean1' : None, 'mean2' : None, \
                               'Sigma1' : None, 'Sigma2' : None, \
                               'nu1': None, 'nu2': None}, seed = None, gaussian = False):
        """
        Multivariate 2-component Student-t mixture random generator.
        Direct random sampling via using the Student-t representation as a continous scale mixture distr.
        -------
        Input:
        -------
        thetas: Component-wise parameters; note that Sigma1,2 are the scale matrices of the
                Wishart priors of the precision matrices of the Student-t's.
        gaussian: boolean, generate from Gaussian mixture if True, otherwise from Student-t
        seed: set seed for rng.
        """
        self.thetas = thetas ; self.seed = seed ; self.gaussian = gaussian
        if self.seed is not None:
            np.random.seed(seed)

    def draw(self, n_samples = 100, k = 2, p = .5):
        """
        Random number generator:
        Input:
        -------
        n_samples: Number of realizations to generate
        k:         Number of features (Dimension of the t-distr.)
        p:         Success probability Bernoulli(p) p.m.f.
        """
        self.n_samples = n_samples ; self.k = k; self.p = p ;
        m = 2                # number of mixture components
        assert (len(self.thetas['mean1']) == k) & (self.thetas['Sigma1'].shape[0] == k), 'Number of dimensions does not match k!'

        if self.gaussian:
            cov1, cov2 = self.thetas['Sigma1'], self.thetas['Sigma2']
        else:
            cov1 = wishart.rvs(df = self.thetas['nu1'], scale = self.thetas['Sigma1'], size=1)
            cov2 = wishart.rvs(df = self.thetas['nu2'], scale = self.thetas['Sigma2'], size=1)

        self.var1 = self.thetas['nu1']/(self.thetas['nu1']-2)*cov1       # variance covariance matrix of first Student-t component
        self.var2 = self.thetas['nu2']/(self.thetas['nu2']-2)*cov2
        self.phi_is = bernoulli.rvs(p = self.p, size = self.n_samples)          # m=2
        Phi = np.tile(self.phi_is, self.k).reshape(self.k,self.n_samples).T              # repeat phi vector to match with random matrix
        rn1 = np.random.multivariate_normal(self.thetas['mean1'], cov1, self.n_samples)
        rn2 = np.random.multivariate_normal(self.thetas['mean2'], cov2, self.n_samples)
        self.sum1 = np.multiply(Phi, rn1)
        self.sum2 = np.multiply(1-Phi, rn2)
        self.x_draws = np.add(self.sum1,self.sum2)
        return self.phi_is, self.x_draws

    def show2D(self, save_plot=False, legend_on = True, **kwargs):
        """
        Make scatter plot for first two dimensions of the random draws
        """
        x_comp1,y_comp1 = self.sum1[:,0], self.sum1[:,1]
        x_comp2,y_comp2 = self.sum2[:,0], self.sum2[:,1]
        fig = plt.figure() ;
        la = plt.scatter(x_comp1, y_comp1, c="blue", **kwargs)
        lb = plt.scatter(x_comp2, y_comp2, c="orange", **kwargs)
        lc = plt.scatter([self.thetas['mean1'][0], self.thetas['mean2'][0]],
                         [self.thetas['mean1'][1],self.thetas['mean2'][1]], c="black", s=6**2, alpha=.5)
        #plt.title("Draws from 2-component \nmultivariate Student-t mixture \n(first two dimensions shown)")
        plt.xlabel(r'$x_{1}$') ; plt.ylabel(r'$x_{2}$')
        if legend_on:
            plt.legend((la, lb), ('Outlier', 'Inlier'),
                            scatterpoints=1, loc='lower right', ncol=3, fontsize=8)
        plt.show() ;
        if save_plot:
            fig.savefig('mixturePlot2D.jpg')
            print("Saved to:", os.getcwd())

    def show3D(self, save_plot=False, legend_on = True, **kwargs):
        """
        Make scatter plot for first three dimensions of the random draws
        """
        fig = plt.figure() ; ax = Axes3D(fig)
        x_comp1,y_comp1, z_comp1 = self.sum1[:,0], self.sum1[:,1], self.sum1[:,2]
        x_comp2,y_comp2, z_comp2 = self.sum2[:,0], self.sum2[:,1], self.sum2[:,2]
        la = ax.scatter(x_comp1, y_comp1, z_comp1, c="blue", **kwargs)
        lb = ax.scatter(x_comp2, y_comp2, z_comp2, c="orange", **kwargs)
        lc = ax.scatter([self.thetas['mean1'][0], self.thetas['mean2'][0]],
                     [self.thetas['mean1'][1],self.thetas['mean2'][1]],
                     [self.thetas['mean1'][2],self.thetas['mean2'][2]], c="black", s=6**2, alpha=.2)

        #plt.title("Draws from 2-component \nmultivariate mixture \n(first three dimensions shown)")
        ax.set_xlabel(r'$x_{1}$') ; ax.set_ylabel(r'$x_{2}$') ;ax.set_zlabel(r'$x_{3}$')
        if legend_on:
            ax.legend((la, lb), ('Outlier', 'Inlier'),
                        scatterpoints=1, loc='lower left', ncol=3, fontsize=8)
        plt.show();
        if save_plot:
            fig.savefig('mixturePlot3D.jpg')
            print("Saved to:", os.getcwd())

# NEW!!!
#-------------------------------------------------------------------------------------
class mvt_tmix:

    def __init__(self, seed = None):
         self.seed = seed

    def draw(self, K = 10, N = 1*10**5, m = 3, gaussian = False):

        """
        Inputs:
        -------
        N: sample size
        K: Dimension of Normal/Student distr.
        m: number of mixture components
        """
        np.random.seed(self.seed)
        self.st0 = np.random.get_state()   # get initial state of RNG
        #np.random.set_state(self.st0)
        print("Drawing from", m,"component mixture distribution.")
        alphas = gamma.rvs(5, size=m)               # shape parameter
        #print(sum(alphas))                              # equivalent sample size
        self.p = dirichlet.rvs(alpha = alphas, size = 1)[0]
        self.phi_is = multinomial.rvs(1, self.p, size=N)       # draw from categorical p.m.f

        self.x_draws = np.zeros((N,K))
        self.hyper_loc, self.hyper_scale, self.thetas, self.var, self.covs, self.rdraws = dict(), dict(), dict(), tuple(), tuple(), tuple()

        for i in range(m):

              self.hyper_loc["mean"+str(i+1)] = norm.rvs(size = 1, loc = 0, scale = 5)
              self.hyper_scale["scale"+str(i+1)] = 1/gamma.rvs(5, size=1)

              self.thetas["mean"+str(i+1)] = norm.rvs(size = K, loc = self.hyper_loc["mean"+str(i+1)],
                          scale = self.hyper_scale["scale"+str(i+1)])
              self.thetas["Sigma"+str(i+1)] = np.eye(K)*(1/gamma.rvs(5, size=K))
              self.thetas["nu"+str(i+1)] = randint.rvs(K+2, K+10, size=1)[0]

              if gaussian:
                 self.covs += (self.thetas['Sigma'+str(i+1)], )
              else:
                 self.covs += (wishart.rvs(df = self.thetas['nu'+str(i+1)], scale = self.thetas['Sigma'+str(i+1)], size=1),)
                 self.var += (self.thetas["nu"+str(i+1)]/(self.thetas["nu"+str(i+1)]-2)*self.covs[i],)       # variance covariance matrix of first Student-t component
              self.rdraws += (np.random.multivariate_normal(self.thetas["mean"+str(i+1)], self.covs[i], N),)

              self.Phi = np.tile(self.phi_is[:,i], K).reshape(K,N).T              # repeat phi vector to match with random matrix
              self.x_draws += np.multiply(self.Phi, self.rdraws[i])

        return self.x_draws, np.argmax(self.phi_is,1)   # X, latent

    def plot(self, draws = None, save_plot=False, legend_on = True, plot_type = ['2D', '3D'], **kwargs):
        """
        Make scatter plot for first two dimensions of the random draws
        """
        if draws is not None:
            self.draws = draws

        #if plot_type[0] == '2D':
        x_comp1,y_comp1 = self.x_draws[:,0], self.x_draws[:,1]
        #x_comp2,y_comp2 = self.sum2[:,0], self.sum2[:,1]
        fig = plt.figure() ;
        la = plt.scatter(x_comp1, y_comp1, c="blue", **kwargs)
        #lb = plt.scatter(x_comp2, y_comp2, c="orange", **kwargs)
        #lc = plt.scatter([self.thetas['mean1'][0], self.thetas['mean2'][0]],
        #                 [self.thetas['mean1'][1],self.thetas['mean2'][1]], c="black", s=6**2, alpha=.5)
        #plt.title("Draws from 2-component \nmultivariate Student-t mixture \n(first two dimensions shown)")
        plt.xlabel(r'$x_{1}$') ; plt.ylabel(r'$x_{2}$')

        if plot_type[0] == '3D':
            fig = plt.figure() ; ax = Axes3D(fig)
            x_comp1, y_comp1, z_comp1 = self.x_draws[:,0], self.x_draws[:,1], self.x_draws[:,2]
            #x_comp2,y_comp2, z_comp2 = self.sum2[:,0], self.sum2[:,1], self.sum2[:,2]
            la = ax.scatter(x_comp1, y_comp1, z_comp1, c="blue", **kwargs)
            #lb = ax.scatter(x_comp2, y_comp2, z_comp2, c="orange", **kwargs)
            #lc = ax.scatter([self.thetas['mean1'][0], self.thetas['mean2'][0]],
            #             [self.thetas['mean1'][1],self.thetas['mean2'][1]],
            #             [self.thetas['mean1'][2],self.thetas['mean2'][2]], c="black", s=6**2, alpha=.2)

            #plt.title("Draws from 2-component \nmultivariate mixture \n(first three dimensions shown)")
            ax.set_xlabel(r'$x_{1}$') ; ax.set_ylabel(r'$x_{2}$') ;ax.set_zlabel(r'$x_{3}$')
            #if legend_on:
            #    ax.legend((la, lb), ('Outlier', 'Inlier'),
            #                scatterpoints=1, loc='lower left', ncol=3, fontsize=8)
        #if legend_on:
        #    plt.legend((la, lb), ('Outlier', 'Inlier'),
        #                    scatterpoints=1, loc='lower right', ncol=3, fontsize=8)
        plt.show() ;
        #if save_plot:
        #    fig.savefig('mixturePlot2D.jpg')
        #    print("Saved to:", os.getcwd())


<reponame>koukyo1994/atmaCup5
import logging

import numpy as np
import pandas as pd

from pathlib import Path
from typing import Any, Dict, Optional, Union, Tuple, List

from scipy.sparse import csr_matrix

class RunningState:
    def __init__(self, config: dict, logger: logging.Logger):
        self.config = config

        self.logger = logger

        self.feature_dir = Path("features")
        self.output_dir = Path("output")

        self.dataframes: Dict[str, pd.DataFrame] = {}
        self.dataframe_roles: Dict[str, str] = {}

        self.callbacks: Dict[str, Dict[str, list]] = {}

        self.data_stats: Dict[str, Optional[Union[str, dict]]] = {}

        self.target_name = ""
        self.target: Optional[Union[pd.DataFrame, Union[pd.Series, np.
                                                        ndarray]]] = None
        self.id_columns: Dict[str, Optional[str]] = {}
        self.connect_to: Dict[str, Optional[str]] = {}
        self.connect_on: Dict[str, Optional[str]] = {}

        self.features: Dict[
            str, Dict[str, Union[pd.DataFrame, Dict[str, csr_matrix]]]] = {}
        self.importances: Dict[str, Union[Dict[str, float], pd.DataFrame]] = {}

        self.splits: List[Tuple[np.ndarray, np.ndarray]] = []

        self.models: Dict[str, Dict[str, Any]] = {}
        self.metrics: Dict[str, Dict[str, float]] = {}

        self.predictions: Dict[str, np.ndarray] = {}

        self.misc: Dict[str, Any] = {}