"""
factor.py
"""
from operator import attrgetter
from numbers import Number

from numpy import float64, inf
from toolz import curry

from zipline.errors import (
    UnknownRankMethod,
    UnsupportedDataType,
)
from zipline.lib.rank import masked_rankdata_2d
from zipline.pipeline.mixins import (
    CustomTermMixin,
    PositiveWindowLengthMixin,
    SingleInputMixin,
)
from zipline.pipeline.term import CompositeTerm, NotSpecified
from zipline.pipeline.expression import (
    BadBinaryOperator,
    COMPARISONS,
    is_comparison,
    MATH_BINOPS,
    method_name_for_op,
    NumericalExpression,
    NUMEXPR_MATH_FUNCS,
    UNARY_OPS,
    unary_op_name,
)
from zipline.pipeline.filters import (
    NumExprFilter,
    PercentileFilter,
)
from zipline.utils.control_flow import nullctx
from zipline.utils.numpy_utils import (
    bool_dtype,
    datetime64ns_dtype,
    float64_dtype,
    int64_dtype,
)
from zipline.utils.preprocess import preprocess

_RANK_METHODS = frozenset(['average', 'min', 'max', 'dense', 'ordinal'])

def numbers_to_float64(func, argname, argvalue):
    """
    Preprocessor for converting numerical inputs into floats.

    This is used in the binary operator constructors for Factor so that
    `2 + Factor()` has the same behavior as `2.0 + Factor()`.
    """
    if isinstance(argvalue, Number):
        return float64(argvalue)
    return argvalue

@curry
def set_attribute(name, value):
    """
    Decorator factory for setting attributes on a function.

    Doesn't change the behavior of the wrapped function.

    Usage
    -----
    >>> @set_attribute('__name__', 'foo')
    ... def bar():
    ...     return 3
    ...
    >>> bar()
    3
    >>> bar.__name__
    'foo'
    """
    def decorator(f):
        setattr(f, name, value)
        return f
    return decorator

# Decorators for setting the __name__ and __doc__ properties of a decorated
# function.
# Example:
with_name = set_attribute('__name__')
with_doc = set_attribute('__doc__')

def binop_return_type(op):
    if is_comparison(op):
        return NumExprFilter
    else:
        return NumExprFactor

def binop_return_dtype(op, left, right):
    """
    Compute the expected return dtype for the given binary operator.

    Parameters
    ----------
    op : str
        Operator symbol, (e.g. '+', '-', ...).
    left : numpy.dtype
        Dtype of left hand side.
    right : numpy.dtype
        Dtype of right hand side.

    Returns
    -------
    outdtype : numpy.dtype
        The dtype of the result of `left <op> right`.
    """
    if is_comparison(op):
        if left != right:
            raise TypeError(
                "Don't know how to compute {left} {op} {right}.\n"
                "Comparisons are only supported between Factors of equal "
                "dtypes.".format(left=left, op=op, right=right)
            )
        return bool_dtype

    elif left != float64_dtype or right != float64_dtype:
        raise TypeError(
            "Don't know how to compute {left} {op} {right}.\n"
            "Arithmetic operators are only supported on Factors of "
            "dtype 'float64'.".format(
                left=left.name,
                op=op,
                right=right.name,
            )
        )
    return float64_dtype

def binary_operator(op):
    """
    Factory function for making binary operator methods on a Factor subclass.

    Returns a function, "binary_operator" suitable for implementing functions
    like __add__.
    """
    # When combining a Factor with a NumericalExpression, we use this
    # attrgetter instance to defer to the commuted implementation of the
    # NumericalExpression operator.
    commuted_method_getter = attrgetter(method_name_for_op(op, commute=True))

    @preprocess(other=numbers_to_float64)
    @with_doc("Binary Operator: '%s'" % op)
    @with_name(method_name_for_op(op))
    def binary_operator(self, other):
        # This can't be hoisted up a scope because the types returned by
        # binop_return_type aren't defined when the top-level function is
        # invoked in the class body of Factor.
        return_type = binop_return_type(op)
        if isinstance(self, NumExprFactor):
            self_expr, other_expr, new_inputs = self.build_binary_op(
                op, other,
            )
            return return_type(
                "({left}) {op} ({right})".format(
                    left=self_expr,
                    op=op,
                    right=other_expr,
                ),
                new_inputs,
                dtype=binop_return_dtype(op, self.dtype, other.dtype),
            )
        elif isinstance(other, NumExprFactor):
            # NumericalExpression overrides ops to correctly handle merging of
            # inputs.  Look up and call the appropriate reflected operator with
            # ourself as the input.
            return commuted_method_getter(other)(self)
        elif isinstance(other, Factor):
            if self is other:
                return return_type(
                    "x_0 {op} x_0".format(op=op),
                    (self,),
                    dtype=binop_return_dtype(op, self.dtype, other.dtype),
                )
            return return_type(
                "x_0 {op} x_1".format(op=op),
                (self, other),
                dtype=binop_return_dtype(op, self.dtype, other.dtype),
            )
        elif isinstance(other, Number):
            return return_type(
                "x_0 {op} ({constant})".format(op=op, constant=other),
                binds=(self,),
                # Interpret numeric literals as floats.
                dtype=binop_return_dtype(op, self.dtype, other.dtype)
            )
        raise BadBinaryOperator(op, self, other)

    return binary_operator

def reflected_binary_operator(op):
    """
    Factory function for making binary operator methods on a Factor.

    Returns a function, "reflected_binary_operator" suitable for implementing
    functions like __radd__.
    """
    assert not is_comparison(op)

    @preprocess(other=numbers_to_float64)
    @with_name(method_name_for_op(op, commute=True))
    def reflected_binary_operator(self, other):

        if isinstance(self, NumericalExpression):
            self_expr, other_expr, new_inputs = self.build_binary_op(
                op, other
            )
            return NumExprFactor(
                "({left}) {op} ({right})".format(
                    left=other_expr,
                    right=self_expr,
                    op=op,
                ),
                new_inputs,
                dtype=binop_return_dtype(op, other.dtype, self.dtype)
            )

        # Only have to handle the numeric case because in all other valid cases
        # the corresponding left-binding method will be called.
        elif isinstance(other, Number):
            return NumExprFactor(
                "{constant} {op} x_0".format(op=op, constant=other),
                binds=(self,),
                dtype=binop_return_dtype(op, other.dtype, self.dtype),
            )
        raise BadBinaryOperator(op, other, self)
    return reflected_binary_operator

def unary_operator(op):
    """
    Factory function for making unary operator methods for Factors.
    """
    # Only negate is currently supported.
    valid_ops = {'-'}
    if op not in valid_ops:
        raise ValueError("Invalid unary operator %s." % op)

    @with_doc("Unary Operator: '%s'" % op)
    @with_name(unary_op_name(op))
    def unary_operator(self):
        if self.dtype != float64_dtype:
            raise TypeError(
                "Can't apply unary operator {op!r} to instance of "
                "{typename!r} with dtype {dtypename!r}.\n"
                "{op!r} is only supported for Factors of dtype "
                "'float64'.".format(
                    op=op,
                    typename=type(self).__name__,
                    dtypename=self.dtype.name,
                )
            )

        # This can't be hoisted up a scope because the types returned by
        # unary_op_return_type aren't defined when the top-level function is
        # invoked.
        if isinstance(self, NumericalExpression):
            return NumExprFactor(
                "{op}({expr})".format(op=op, expr=self._expr),
                self.inputs,
                dtype=float64_dtype,
            )
        else:
            return NumExprFactor(
                "{op}x_0".format(op=op),
                (self,),
                dtype=float64_dtype,
            )
    return unary_operator

def function_application(func):
    """
    Factory function for producing function application methods for Factor
    subclasses.
    """
    if func not in NUMEXPR_MATH_FUNCS:
        raise ValueError("Unsupported mathematical function '%s'" % func)

    @with_name(func)
    def mathfunc(self):
        if isinstance(self, NumericalExpression):
            return NumExprFactor(
                "{func}({expr})".format(func=func, expr=self._expr),
                self.inputs,
                dtype=float64_dtype,
            )
        else:
            return NumExprFactor(
                "{func}(x_0)".format(func=func),
                (self,),
                dtype=float64_dtype,
            )
    return mathfunc

FACTOR_DTYPES = frozenset([datetime64ns_dtype, float64_dtype, int64_dtype])

class Factor(CompositeTerm):
    """
    Pipeline API expression producing numerically-valued outputs.
    """
    # Dynamically add functions for creating NumExprFactor/NumExprFilter
    # instances.
    clsdict = locals()
    clsdict.update(
        {
            method_name_for_op(op): binary_operator(op)
            # Don't override __eq__ because it breaks comparisons on tuples of
            # Factors.
            for op in MATH_BINOPS.union(COMPARISONS - {'=='})
        }
    )
    clsdict.update(
        {
            method_name_for_op(op, commute=True): reflected_binary_operator(op)
            for op in MATH_BINOPS
        }
    )
    clsdict.update(
        {
            unary_op_name(op): unary_operator(op)
            for op in UNARY_OPS
        }
    )

    clsdict.update(
        {
            funcname: function_application(funcname)
            for funcname in NUMEXPR_MATH_FUNCS
        }
    )

    __truediv__ = clsdict['__div__']
    __rtruediv__ = clsdict['__rdiv__']

    eq = binary_operator('==')

    def _validate(self):
        # Do superclass validation first so that `NotSpecified` dtypes get
        # handled.
        retval = super(Factor, self)._validate()
        if self.dtype not in FACTOR_DTYPES:
            raise UnsupportedDataType(
                typename=type(self).__name__,
                dtype=self.dtype
            )
        return retval

    def rank(self, method='ordinal', ascending=True, mask=NotSpecified):
        """
        Construct a new Factor representing the sorted rank of each column
        within each row.

        Parameters
        ----------
        method : str, {'ordinal', 'min', 'max', 'dense', 'average'}
            The method used to assign ranks to tied elements. See
            `scipy.stats.rankdata` for a full description of the semantics for
            each ranking method. Default is 'ordinal'.
        ascending : bool, optional
            Whether to return sorted rank in ascending or descending order.
            Default is True.
        mask : zipline.pipeline.Filter, optional
            A Filter representing assets to consider when computing ranks.
            If mask is supplied, ranks are computed ignoring any asset/date
            pairs for which `mask` produces a value of False.

        Returns
        -------
        ranks : zipline.pipeline.factors.Rank
            A new factor that will compute the ranking of the data produced by
            `self`.

        Notes
        -----
        The default value for `method` is different from the default for
        `scipy.stats.rankdata`.  See that function's documentation for a full
        description of the valid inputs to `method`.

        Missing or non-existent data on a given day will cause an asset to be
        given a rank of NaN for that day.

        See Also
        --------
        scipy.stats.rankdata
        zipline.lib.rank.masked_rankdata_2d
        zipline.pipeline.factors.factor.Rank
        """
        return Rank(self, method=method, ascending=ascending, mask=mask)

    def top(self, N, mask=NotSpecified):
        """
        Construct a Filter matching the top N asset values of self each day.

        Parameters
        ----------
        N : int
            Number of assets passing the returned filter each day.
        mask : zipline.pipeline.Filter, optional
            A Filter representing assets to consider when computing ranks.
            If mask is supplied, top values are computed ignoring any
            asset/date pairs for which `mask` produces a value of False.

        Returns
        -------
        filter : zipline.pipeline.filters.Filter
        """
        return self.rank(ascending=False, mask=mask) <= N

    def bottom(self, N, mask=NotSpecified):
        """
        Construct a Filter matching the bottom N asset values of self each day.

        Parameters
        ----------
        N : int
            Number of assets passing the returned filter each day.
        mask : zipline.pipeline.Filter, optional
            A Filter representing assets to consider when computing ranks.
            If mask is supplied, bottom values are computed ignoring any
            asset/date pairs for which `mask` produces a value of False.

        Returns
        -------
        filter : zipline.pipeline.Filter
        """
        return self.rank(ascending=True, mask=mask) <= N

    def percentile_between(self,
                           min_percentile,
                           max_percentile,
                           mask=NotSpecified):
        """
        Construct a new Filter representing entries from the output of this
        Factor that fall within the percentile range defined by min_percentile
        and max_percentile.

        Parameters
        ----------
        min_percentile : float [0.0, 100.0]
            Return True for assets falling above this percentile in the data.
        max_percentile : float [0.0, 100.0]
            Return True for assets falling below this percentile in the data.
        mask : zipline.pipeline.Filter, optional
            A Filter representing assets to consider when percentile
            calculating thresholds.  If mask is supplied, percentile cutoffs
            are computed each day using only assets for which ``mask`` returns
            True.  Assets for which ``mask`` produces False will produce False
            in the output of this Factor as well.

        Returns
        -------
        out : zipline.pipeline.filters.PercentileFilter
            A new filter that will compute the specified percentile-range mask.

        See Also
        --------
        zipline.pipeline.filters.filter.PercentileFilter
        """
        return PercentileFilter(
            self,
            min_percentile=min_percentile,
            max_percentile=max_percentile,
            mask=mask,
        )

    def isnan(self):
        """
        A Filter producing True for all values where this Factor is NaN.

        Returns
        -------
        nanfilter : zipline.pipeline.filters.Filter
        """
        return self != self

    def notnan(self):
        """
        A Filter producing True for values where this Factor is not NaN.

        Returns
        -------
        nanfilter : zipline.pipeline.filters.Filter
        """
        return ~self.isnan()

    def isfinite(self):
        """
        A Filter producing True for values where this Factor is anything but
        NaN, inf, or -inf.
        """
        return (-inf < self) & (self < inf)

class NumExprFactor(NumericalExpression, Factor):
    """
    Factor computed from a numexpr expression.

    Parameters
    ----------
    expr : string
       A string suitable for passing to numexpr.  All variables in 'expr'
       should be of the form "x_i", where i is the index of the corresponding
       factor input in 'binds'.
    binds : tuple
       A tuple of factors to use as inputs.

    Notes
    -----
    NumExprFactors are constructed by numerical operators like `+` and `-`.
    Users should rarely need to construct a NumExprFactor directly.
    """
    pass

class Rank(SingleInputMixin, Factor):
    """
    A Factor representing the row-wise rank data of another Factor.

    Parameters
    ----------
    factor : zipline.pipeline.factors.Factor
        The factor on which to compute ranks.
    method : str, {'average', 'min', 'max', 'dense', 'ordinal'}
        The method used to assign ranks to tied elements.  See
        `scipy.stats.rankdata` for a full description of the semantics for each
        ranking method.

    See Also
    --------
    scipy.stats.rankdata : Underlying ranking algorithm.
    zipline.factors.Factor.rank : Method-style interface to same functionality.

    Notes
    -----
    Most users should call Factor.rank rather than directly construct an
    instance of this class.
    """
    window_length = 0
    dtype = float64_dtype

    def __new__(cls, factor, method, ascending, mask):
        return super(Rank, cls).__new__(
            cls,
            inputs=(factor,),
            method=method,
            ascending=ascending,
            mask=mask,
        )

    def _init(self, method, ascending, *args, **kwargs):
        self._method = method
        self._ascending = ascending
        return super(Rank, self)._init(*args, **kwargs)

    @classmethod
    def static_identity(cls, method, ascending, *args, **kwargs):
        return (
            super(Rank, cls).static_identity(*args, **kwargs),
            method,
            ascending,
        )

    def _validate(self):
        """
        Verify that the stored rank method is valid.
        """
        if self._method not in _RANK_METHODS:
            raise UnknownRankMethod(
                method=self._method,
                choices=set(_RANK_METHODS),
            )
        return super(Rank, self)._validate()

    def _compute(self, arrays, dates, assets, mask):
        """
        For each row in the input, compute a like-shaped array of per-row
        ranks.
        """
        return masked_rankdata_2d(
            arrays[0],
            mask,
            self.inputs[0].missing_value,
            self._method,
            self._ascending,
        )

    def __repr__(self):
        return "{type}({input_}, method='{method}', mask={mask})".format(
            type=type(self).__name__,
            input_=self.inputs[0],
            method=self._method,
            mask=self.mask,
        )

class CustomFactor(PositiveWindowLengthMixin, CustomTermMixin, Factor):
    '''
    Base class for user-defined Factors.

    Parameters
    ----------
    inputs : iterable, optional
        An iterable of `BoundColumn` instances (e.g. USEquityPricing.close),
        describing the data to load and pass to `self.compute`.  If this
        argument is passed to the CustomFactor constructor, we look for a
        class-level attribute named `inputs`.
    window_length : int, optional
        Number of rows to pass for each input.  If this argument is not passed
        to the CustomFactor constructor, we look for a class-level attribute
        named `window_length`.

    Notes
    -----
    Users implementing their own Factors should subclass CustomFactor and
    implement a method named `compute` with the following signature:

    .. code-block:: python

        def compute(self, today, assets, out, *inputs):
           ...

    On each simulation date, ``compute`` will be called with the current date,
    an array of sids, an output array, and an input array for each expression
    passed as inputs to the CustomFactor constructor.

    The specific types of the values passed to `compute` are as follows::

        today : np.datetime64[ns]
            Row label for the last row of all arrays passed as `inputs`.
        assets : np.array[int64, ndim=1]
            Column labels for `out` and`inputs`.
        out : np.array[self.dtype, ndim=1]
            Output array of the same shape as `assets`.  `compute` should write
            its desired return values into `out`.
        *inputs : tuple of np.array
            Raw data arrays corresponding to the values of `self.inputs`.

    ``compute`` functions should expect to be passed NaN values for dates on
    which no data was available for an asset.  This may include dates on which
    an asset did not yet exist.

    For example, if a CustomFactor requires 10 rows of close price data, and
    asset A started trading on Monday June 2nd, 2014, then on Tuesday, June
    3rd, 2014, the column of input data for asset A will have 9 leading NaNs
    for the preceding days on which data was not yet available.

    Examples
    --------

    A CustomFactor with pre-declared defaults:

    .. code-block:: python

        class TenDayRange(CustomFactor):
            """
            Computes the difference between the highest high in the last 10
            days and the lowest low.

            Pre-declares high and low as default inputs and `window_length` as
            10.
            """

            inputs = [USEquityPricing.high, USEquityPricing.low]
            window_length = 10

            def compute(self, today, assets, out, highs, lows):
                from numpy import nanmin, nanmax

                highest_highs = nanmax(highs, axis=0)
                lowest_lows = nanmin(lows, axis=0)
                out[:] = highest_highs - lowest_lows

        # Doesn't require passing inputs or window_length because they're
        # pre-declared as defaults for the TenDayRange class.
        ten_day_range = TenDayRange()

    A CustomFactor without defaults:

    .. code-block:: python

        class MedianValue(CustomFactor):
            """
            Computes the median value of an arbitrary single input over an
            arbitrary window..

            Does not declare any defaults, so values for `window_length` and
            `inputs` must be passed explicitly on every construction.
            """

            def compute(self, today, assets, out, data):
                from numpy import nanmedian
                out[:] = data.nanmedian(data, axis=0)

        # Values for `inputs` and `window_length` must be passed explicitly to
        # MedianValue.
        median_close10 = MedianValue([USEquityPricing.close], window_length=10)
        median_low15 = MedianValue([USEquityPricing.low], window_length=15)
    '''
    dtype = float64_dtype
    ctx = nullctx()


import collections
import ctypes
import io
import math
import numpy as np
import os
import pyogg
import tempfile
import wave

from .helpers import LimitingPool
from collections import namedtuple
from .io import open_remote, remove_remote, copy_remote, is_remote_path

AudioFormat = namedtuple('AudioFormat', 'rate channels width')

DEFAULT_RATE = 16000
DEFAULT_CHANNELS = 1
DEFAULT_WIDTH = 2
DEFAULT_FORMAT = AudioFormat(DEFAULT_RATE, DEFAULT_CHANNELS, DEFAULT_WIDTH)

AUDIO_TYPE_NP = 'application/vnd.mozilla.np'
AUDIO_TYPE_PCM = 'application/vnd.mozilla.pcm'
AUDIO_TYPE_WAV = 'audio/wav'
AUDIO_TYPE_OPUS = 'application/vnd.mozilla.opus'
AUDIO_TYPE_OGG_OPUS = 'application/vnd.deepspeech.ogg_opus'

SERIALIZABLE_AUDIO_TYPES = [AUDIO_TYPE_WAV, AUDIO_TYPE_OPUS, AUDIO_TYPE_OGG_OPUS]

OPUS_PCM_LEN_SIZE = 4
OPUS_RATE_SIZE = 4
OPUS_CHANNELS_SIZE = 1
OPUS_WIDTH_SIZE = 1
OPUS_CHUNK_LEN_SIZE = 2

class Sample:
    """
    Represents in-memory audio data of a certain (convertible) representation.

    Attributes
    ----------
    audio_type : str
        See `__init__`.
    audio_format : util.audio.AudioFormat
        See `__init__`.
    audio : binary
        Audio data represented as indicated by `audio_type`
    duration : float
        Audio duration of the sample in seconds
    """
    def __init__(self, audio_type, raw_data, audio_format=None, sample_id=None):
        """
        Parameters
        ----------
        audio_type : str
            Audio data representation type
            Supported types:
                - util.audio.AUDIO_TYPE_OPUS: Memory file representation (BytesIO) of Opus encoded audio
                    wrapped by a custom container format (used in SDBs)
                - util.audio.AUDIO_TYPE_WAV: Memory file representation (BytesIO) of a Wave file
                - util.audio.AUDIO_TYPE_PCM: Binary representation (bytearray) of PCM encoded audio data (Wave file without header)
                - util.audio.AUDIO_TYPE_NP: NumPy representation of audio data (np.float32) - typically used for GPU feeding
        raw_data : binary
            Audio data in the form of the provided representation type (see audio_type).
            For types util.audio.AUDIO_TYPE_OPUS or util.audio.AUDIO_TYPE_WAV data can also be passed as a bytearray.
        audio_format : util.audio.AudioFormat
            Required in case of audio_type = util.audio.AUDIO_TYPE_PCM or util.audio.AUDIO_TYPE_NP,
            as this information cannot be derived from raw audio data.
        sample_id : str
            Tracking ID - should indicate sample's origin as precisely as possible
        """
        self.audio_type = audio_type
        self.audio_format = audio_format
        self.sample_id = sample_id
        if audio_type in SERIALIZABLE_AUDIO_TYPES:
            self.audio = raw_data if isinstance(raw_data, io.BytesIO) else io.BytesIO(raw_data)
            self.duration = read_duration(audio_type, self.audio)
            if not self.audio_format:
                self.audio_format = read_format(audio_type, self.audio)
        else:
            self.audio = raw_data
            if self.audio_format is None:
                raise ValueError('For audio type "{}" parameter "audio_format" is mandatory'.format(self.audio_type))
            if audio_type == AUDIO_TYPE_PCM:
                self.duration = get_pcm_duration(len(self.audio), self.audio_format)
            elif audio_type == AUDIO_TYPE_NP:
                self.duration = get_np_duration(len(self.audio), self.audio_format)
            else:
                raise ValueError('Unsupported audio type: {}'.format(self.audio_type))

    def change_audio_type(self, new_audio_type, bitrate=None):
        """
        In-place conversion of audio data into a different representation.

        Parameters
        ----------
        new_audio_type : str
            New audio-type - see `__init__`.
        bitrate : int
            Bitrate to use in case of converting to a lossy audio-type.
        """
        if self.audio_type == new_audio_type:
            return
        if new_audio_type == AUDIO_TYPE_PCM and self.audio_type in SERIALIZABLE_AUDIO_TYPES:
            self.audio_format, audio = read_audio(self.audio_type, self.audio)
            self.audio.close()
            self.audio = audio
        elif new_audio_type == AUDIO_TYPE_PCM and self.audio_type == AUDIO_TYPE_NP:
            self.audio = np_to_pcm(self.audio, self.audio_format)
        elif new_audio_type == AUDIO_TYPE_NP:
            self.change_audio_type(AUDIO_TYPE_PCM)
            self.audio = pcm_to_np(self.audio, self.audio_format)
        elif new_audio_type in SERIALIZABLE_AUDIO_TYPES:
            self.change_audio_type(AUDIO_TYPE_PCM)
            audio_bytes = io.BytesIO()
            write_audio(new_audio_type, audio_bytes, self.audio, audio_format=self.audio_format, bitrate=bitrate)
            audio_bytes.seek(0)
            self.audio = audio_bytes
        else:
            raise RuntimeError('Changing audio representation type from "{}" to "{}" not supported'
                               .format(self.audio_type, new_audio_type))
        self.audio_type = new_audio_type

def _unpack_and_change_audio_type(sample_and_audio_type):
    packed_sample, audio_type, bitrate = sample_and_audio_type
    if hasattr(packed_sample, 'unpack'):
        sample = packed_sample.unpack()
    else:
        sample = packed_sample
    sample.change_audio_type(audio_type, bitrate=bitrate)
    return sample

def change_audio_types(packed_samples, audio_type=AUDIO_TYPE_PCM, bitrate=None, processes=None, process_ahead=None):
    with LimitingPool(processes=processes, process_ahead=process_ahead) as pool:
        yield from pool.imap(_unpack_and_change_audio_type, map(lambda s: (s, audio_type, bitrate), packed_samples))

def get_loadable_audio_type_from_extension(ext):
    return {
        '.wav': AUDIO_TYPE_WAV,
        '.opus': AUDIO_TYPE_OGG_OPUS,
    }.get(ext, None)

def read_audio_format_from_wav_file(wav_file):
    return AudioFormat(wav_file.getframerate(), wav_file.getnchannels(), wav_file.getsampwidth())

def get_num_samples(pcm_buffer_size, audio_format=DEFAULT_FORMAT):
    return pcm_buffer_size // (audio_format.channels * audio_format.width)

def get_pcm_duration(pcm_buffer_size, audio_format=DEFAULT_FORMAT):
    """Calculates duration in seconds of a binary PCM buffer (typically read from a WAV file)"""
    return get_num_samples(pcm_buffer_size, audio_format) / audio_format.rate

def get_np_duration(np_len, audio_format=DEFAULT_FORMAT):
    """Calculates duration in seconds of NumPy audio data"""
    return np_len / audio_format.rate

def convert_audio(src_audio_path, dst_audio_path, file_type=None, audio_format=DEFAULT_FORMAT):
    import sox
    transformer = sox.Transformer()
    transformer.set_output_format(file_type=file_type,
                                  rate=audio_format.rate,
                                  channels=audio_format.channels,
                                  bits=audio_format.width * 8)
    transformer.build(src_audio_path, dst_audio_path)

class AudioFile:
    """
    Audio data file wrapper that ensures that the file is loaded with the correct sample rate, channels,
    and width, and converts the file on the fly otherwise.
    """
    def __init__(self, audio_path, as_path=False, audio_format=DEFAULT_FORMAT):
        self.audio_path = audio_path
        self.audio_format = audio_format
        self.as_path = as_path
        self.open_file = None
        self.open_wav = None
        self.tmp_file_path = None
        self.tmp_src_file_path = None

    def __enter__(self):
        if self.audio_path.endswith('.wav'):
            self.open_file = open_remote(self.audio_path, 'rb')
            self.open_wav = wave.open(self.open_file)
            if read_audio_format_from_wav_file(self.open_wav) == self.audio_format:
                if self.as_path:
                    self.open_wav.close()
                    self.open_file.close()
                    return self.audio_path
                return self.open_wav
            self.open_wav.close()
            self.open_file.close()

        # If the format isn't right, copy the file to local tmp dir and do the conversion on disk
        if is_remote_path(self.audio_path):
            _, self.tmp_src_file_path = tempfile.mkstemp(suffix='.wav')
            copy_remote(self.audio_path, self.tmp_src_file_path, True)
            self.audio_path = self.tmp_src_file_path

        _, self.tmp_file_path = tempfile.mkstemp(suffix='.wav')
        convert_audio(self.audio_path, self.tmp_file_path, file_type='wav', audio_format=self.audio_format)
        if self.as_path:
            return self.tmp_file_path
        self.open_wav = wave.open(self.tmp_file_path, 'rb')
        return self.open_wav

    def __exit__(self, *args):
        if not self.as_path:
            self.open_wav.close()
            if self.open_file:
                self.open_file.close()
        if self.tmp_file_path is not None:
            os.remove(self.tmp_file_path)
        if self.tmp_src_file_path is not None:
            os.remove(self.tmp_src_file_path)

def read_frames(wav_file, frame_duration_ms=30, yield_remainder=False):
    audio_format = read_audio_format_from_wav_file(wav_file)
    frame_size = int(audio_format.rate * (frame_duration_ms / 1000.0))
    while True:
        try:
            data = wav_file.readframes(frame_size)
            if not yield_remainder and get_pcm_duration(len(data), audio_format) * 1000 < frame_duration_ms:
                break
            yield data
        except EOFError:
            break

def read_frames_from_file(audio_path, audio_format=DEFAULT_FORMAT, frame_duration_ms=30, yield_remainder=False):
    with AudioFile(audio_path, audio_format=audio_format) as wav_file:
        for frame in read_frames(wav_file, frame_duration_ms=frame_duration_ms, yield_remainder=yield_remainder):
            yield frame

def vad_split(audio_frames,
              audio_format=DEFAULT_FORMAT,
              num_padding_frames=10,
              threshold=0.5,
              aggressiveness=3):
    from webrtcvad import Vad  # pylint: disable=import-outside-toplevel
    if audio_format.channels != 1:
        raise ValueError('VAD-splitting requires mono samples')
    if audio_format.width != 2:
        raise ValueError('VAD-splitting requires 16 bit samples')
    if audio_format.rate not in [8000, 16000, 32000, 48000]:
        raise ValueError('VAD-splitting only supported for sample rates 8000, 16000, 32000, or 48000')
    if aggressiveness not in [0, 1, 2, 3]:
        raise ValueError('VAD-splitting aggressiveness mode has to be one of 0, 1, 2, or 3')
    ring_buffer = collections.deque(maxlen=num_padding_frames)
    triggered = False
    vad = Vad(int(aggressiveness))
    voiced_frames = []
    frame_duration_ms = 0
    frame_index = 0
    for frame_index, frame in enumerate(audio_frames):
        frame_duration_ms = get_pcm_duration(len(frame), audio_format) * 1000
        if int(frame_duration_ms) not in [10, 20, 30]:
            raise ValueError('VAD-splitting only supported for frame durations 10, 20, or 30 ms')
        is_speech = vad.is_speech(frame, audio_format.rate)
        if not triggered:
            ring_buffer.append((frame, is_speech))
            num_voiced = len([f for f, speech in ring_buffer if speech])
            if num_voiced > threshold * ring_buffer.maxlen:
                triggered = True
                for f, s in ring_buffer:
                    voiced_frames.append(f)
                ring_buffer.clear()
        else:
            voiced_frames.append(frame)
            ring_buffer.append((frame, is_speech))
            num_unvoiced = len([f for f, speech in ring_buffer if not speech])
            if num_unvoiced > threshold * ring_buffer.maxlen:
                triggered = False
                yield b''.join(voiced_frames), \
                      frame_duration_ms * max(0, frame_index - len(voiced_frames)), \
                      frame_duration_ms * frame_index
                ring_buffer.clear()
                voiced_frames = []
    if len(voiced_frames) > 0:
        yield b''.join(voiced_frames), \
              frame_duration_ms * (frame_index - len(voiced_frames)), \
              frame_duration_ms * (frame_index + 1)

def pack_number(n, num_bytes):
    return n.to_bytes(num_bytes, 'big', signed=False)

def unpack_number(data):
    return int.from_bytes(data, 'big', signed=False)

def get_opus_frame_size(rate):
    return 60 * rate // 1000

def write_opus(opus_file, audio_data, audio_format=DEFAULT_FORMAT, bitrate=None):
    frame_size = get_opus_frame_size(audio_format.rate)
    import opuslib  # pylint: disable=import-outside-toplevel
    encoder = opuslib.Encoder(audio_format.rate, audio_format.channels, 'audio')
    if bitrate is not None:
        encoder.bitrate = bitrate
    chunk_size = frame_size * audio_format.channels * audio_format.width
    opus_file.write(pack_number(len(audio_data), OPUS_PCM_LEN_SIZE))
    opus_file.write(pack_number(audio_format.rate, OPUS_RATE_SIZE))
    opus_file.write(pack_number(audio_format.channels, OPUS_CHANNELS_SIZE))
    opus_file.write(pack_number(audio_format.width, OPUS_WIDTH_SIZE))
    for i in range(0, len(audio_data), chunk_size):
        chunk = audio_data[i:i + chunk_size]
        # Preventing non-deterministic encoding results from uninitialized remainder of the encoder buffer
        if len(chunk) < chunk_size:
            chunk = chunk + b'\0' * (chunk_size - len(chunk))
        encoded = encoder.encode(chunk, frame_size)
        opus_file.write(pack_number(len(encoded), OPUS_CHUNK_LEN_SIZE))
        opus_file.write(encoded)

def read_opus_header(opus_file):
    opus_file.seek(0)
    pcm_buffer_size = unpack_number(opus_file.read(OPUS_PCM_LEN_SIZE))
    rate = unpack_number(opus_file.read(OPUS_RATE_SIZE))
    channels = unpack_number(opus_file.read(OPUS_CHANNELS_SIZE))
    width = unpack_number(opus_file.read(OPUS_WIDTH_SIZE))
    return pcm_buffer_size, AudioFormat(rate, channels, width)

def read_opus(opus_file):
    pcm_buffer_size, audio_format = read_opus_header(opus_file)
    frame_size = get_opus_frame_size(audio_format.rate)
    import opuslib  # pylint: disable=import-outside-toplevel
    decoder = opuslib.Decoder(audio_format.rate, audio_format.channels)
    audio_data = bytearray()
    while len(audio_data) < pcm_buffer_size:
        chunk_len = unpack_number(opus_file.read(OPUS_CHUNK_LEN_SIZE))
        chunk = opus_file.read(chunk_len)
        decoded = decoder.decode(chunk, frame_size)
        audio_data.extend(decoded)
    audio_data = audio_data[:pcm_buffer_size]
    return audio_format, bytes(audio_data)

def read_ogg_opus(ogg_file):
    error = ctypes.c_int()
    ogg_file_buffer = ogg_file.getbuffer()
    ubyte_array = ctypes.c_ubyte * len(ogg_file_buffer)
    opusfile = pyogg.opus.op_open_memory(
        ubyte_array.from_buffer(ogg_file_buffer),
        len(ogg_file_buffer),
        ctypes.pointer(error)
    )

    if error.value != 0:
        raise ValueError(
            ("Ogg/Opus buffer could not be read."
             "Error code: {}").format(error.value)
        )

    channel_count = pyogg.opus.op_channel_count(opusfile, -1)
    sample_rate = 48000 # opus files are always 48kHz
    sample_width = 2 # always 16-bit
    audio_format = AudioFormat(sample_rate, channel_count, sample_width)

    # Allocate sufficient memory to store the entire PCM
    pcm_size = pyogg.opus.op_pcm_total(opusfile, -1)
    Buf = pyogg.opus.opus_int16*(pcm_size*channel_count)
    buf = Buf()

    # Create a pointer to the newly allocated memory.  It
    # seems we can only do pointer arithmetic on void
    # pointers.  See
    # https://mattgwwalker.wordpress.com/2020/05/30/pointer-manipulation-in-python/
    buf_ptr = ctypes.cast(
        ctypes.pointer(buf),
        ctypes.c_void_p
    )
    assert buf_ptr.value is not None # for mypy
    buf_ptr_zero = buf_ptr.value

    #: Bytes per sample
    bytes_per_sample = ctypes.sizeof(pyogg.opus.opus_int16)

    # Read through the entire file, copying the PCM into the
    # buffer
    samples = 0
    while True:
        # Calculate remaining buffer size
        remaining_buffer = (
            len(buf) # int
            - (buf_ptr.value - buf_ptr_zero) // bytes_per_sample
        )

        # Convert buffer pointer to the desired type
        ptr = ctypes.cast(
            buf_ptr,
            ctypes.POINTER(pyogg.opus.opus_int16)
        )

        # Read the next section of PCM
        ns = pyogg.opus.op_read(
            opusfile,
            ptr,
            remaining_buffer,
            pyogg.ogg.c_int_p()
        )

        # Check for errors
        if ns < 0:
            raise ValueError(
                "Error while reading OggOpus buffer. "+
                "Error code: {}".format(ns)
            )

        # Increment the pointer
        buf_ptr.value += (
            ns
            * bytes_per_sample
            * channel_count
        )
        assert buf_ptr.value is not None # for mypy

        samples += ns

        # Check if we've finished
        if ns == 0:
            break

    # Close the open file
    pyogg.opus.op_free(opusfile)

    # Cast buffer to a one-dimensional array of chars
    #: Raw PCM data from audio file.
    CharBuffer = ctypes.c_byte * (bytes_per_sample * channel_count * pcm_size)
    audio_data = CharBuffer.from_buffer(buf)

    return audio_format, audio_data

def write_wav(wav_file, pcm_data, audio_format=DEFAULT_FORMAT):
    # wav_file is already a file-pointer here
    with wave.open(wav_file, 'wb') as wav_file_writer:
        wav_file_writer.setframerate(audio_format.rate)
        wav_file_writer.setnchannels(audio_format.channels)
        wav_file_writer.setsampwidth(audio_format.width)
        wav_file_writer.writeframes(pcm_data)

def read_wav(wav_file):
    wav_file.seek(0)
    with wave.open(wav_file, 'rb') as wav_file_reader:
        audio_format = read_audio_format_from_wav_file(wav_file_reader)
        pcm_data = wav_file_reader.readframes(wav_file_reader.getnframes())
        return audio_format, pcm_data

def read_audio(audio_type, audio_file):
    if audio_type == AUDIO_TYPE_WAV:
        return read_wav(audio_file)
    if audio_type == AUDIO_TYPE_OPUS:
        return read_opus(audio_file)
    if audio_type == AUDIO_TYPE_OGG_OPUS:
        return read_ogg_opus(audio_file)
    raise ValueError('Unsupported audio type: {}'.format(audio_type))

def write_audio(audio_type, audio_file, pcm_data, audio_format=DEFAULT_FORMAT, bitrate=None):
    if audio_type == AUDIO_TYPE_WAV:
        return write_wav(audio_file, pcm_data, audio_format=audio_format)
    if audio_type == AUDIO_TYPE_OPUS:
        return write_opus(audio_file, pcm_data, audio_format=audio_format, bitrate=bitrate)
    raise ValueError('Unsupported audio type: {}'.format(audio_type))

def read_wav_duration(wav_file):
    wav_file.seek(0)
    with wave.open(wav_file, 'rb') as wav_file_reader:
        return wav_file_reader.getnframes() / wav_file_reader.getframerate()

def read_opus_duration(opus_file):
    pcm_buffer_size, audio_format = read_opus_header(opus_file)
    return get_pcm_duration(pcm_buffer_size, audio_format)

def read_ogg_opus_duration(ogg_file):
    error = ctypes.c_int()
    ogg_file_buffer = ogg_file.getbuffer()
    ubyte_array = ctypes.c_ubyte * len(ogg_file_buffer)
    opusfile = pyogg.opus.op_open_memory(
        ubyte_array.from_buffer(ogg_file_buffer),
        len(ogg_file_buffer),
        ctypes.pointer(error)
    )

    if error.value != 0:
        raise ValueError(
            ("Ogg/Opus buffer could not be read."
             "Error code: {}").format(error.value)
        )

    pcm_buffer_size = pyogg.opus.op_pcm_total(opusfile, -1)
    channel_count = pyogg.opus.op_channel_count(opusfile, -1)
    sample_rate = 48000 # opus files are always 48kHz
    sample_width = 2 # always 16-bit
    audio_format = AudioFormat(sample_rate, channel_count, sample_width)
    pyogg.opus.op_free(opusfile)
    return get_pcm_duration(pcm_buffer_size, audio_format)

def read_duration(audio_type, audio_file):
    if audio_type == AUDIO_TYPE_WAV:
        return read_wav_duration(audio_file)
    if audio_type == AUDIO_TYPE_OPUS:
        return read_opus_duration(audio_file)
    if audio_type == AUDIO_TYPE_OGG_OPUS:
        return read_ogg_opus_duration(audio_file)
    raise ValueError('Unsupported audio type: {}'.format(audio_type))

def read_wav_format(wav_file):
    wav_file.seek(0)
    with wave.open(wav_file, 'rb') as wav_file_reader:
        return read_audio_format_from_wav_file(wav_file_reader)

def read_opus_format(opus_file):
    _, audio_format = read_opus_header(opus_file)
    return audio_format

def read_ogg_opus_format(ogg_file):
    error = ctypes.c_int()
    ogg_file_buffer = ogg_file.getbuffer()
    ubyte_array = ctypes.c_ubyte * len(ogg_file_buffer)
    opusfile = pyogg.opus.op_open_memory(
        ubyte_array.from_buffer(ogg_file_buffer),
        len(ogg_file_buffer),
        ctypes.pointer(error)
    )

    if error.value != 0:
        raise ValueError(
            ("Ogg/Opus buffer could not be read."
             "Error code: {}").format(error.value)
        )

    channel_count = pyogg.opus.op_channel_count(opusfile, -1)
    pyogg.opus.op_free(opusfile)

    sample_rate = 48000 # opus files are always 48kHz
    sample_width = 2 # always 16-bit
    return AudioFormat(sample_rate, channel_count, sample_width)

def read_format(audio_type, audio_file):
    if audio_type == AUDIO_TYPE_WAV:
        return read_wav_format(audio_file)
    if audio_type == AUDIO_TYPE_OPUS:
        return read_opus_format(audio_file)
    if audio_type == AUDIO_TYPE_OGG_OPUS:
        return read_ogg_opus_format(audio_file)
    raise ValueError('Unsupported audio type: {}'.format(audio_type))

def get_dtype(audio_format):
    if audio_format.width not in [1, 2, 4]:
        raise ValueError('Unsupported sample width: {}'.format(audio_format.width))
    return [None, np.int8, np.int16, None, np.int32][audio_format.width]

def pcm_to_np(pcm_data, audio_format=DEFAULT_FORMAT):
    if audio_format.channels != 1:
        raise ValueError('Mono-channel audio required')
    dtype = get_dtype(audio_format)
    samples = np.frombuffer(pcm_data, dtype=dtype)
    samples = samples.astype(np.float32) / np.iinfo(dtype).max
    return np.expand_dims(samples, axis=1)

def np_to_pcm(np_data, audio_format=DEFAULT_FORMAT):
    if audio_format.channels != 1:
        raise ValueError('Mono-channel audio required')
    dtype = get_dtype(audio_format)
    np_data = np_data.squeeze()
    np_data = np_data * np.iinfo(dtype).max
    np_data = np_data.astype(dtype)
    return np_data.tobytes()

def rms_to_dbfs(rms):
    return 20.0 * math.log10(max(1e-16, rms)) + 3.0103

def max_dbfs(sample_data):
    # Peak dBFS based on the maximum energy sample. Will prevent overdrive if used for normalization.
    return rms_to_dbfs(max(abs(np.min(sample_data)), abs(np.max(sample_data))))

def mean_dbfs(sample_data):
    return rms_to_dbfs(math.sqrt(np.mean(np.square(sample_data, dtype=np.float64))))

def gain_db_to_ratio(gain_db):
    return math.pow(10.0, gain_db / 20.0)

def normalize_audio(sample_data, dbfs=3.0103):
    return np.maximum(np.minimum(sample_data * gain_db_to_ratio(dbfs - max_dbfs(sample_data)), 1.0), -1.0)


import os
import argparse
import numpy as np
import tensorflow as tf
import time
from dataloader.data_generator import load_data
from models.cliquenet import build_model

def into_batch(data, label, batch_size, shuffle):
    if shuffle:
        rand_indexes = np.random.permutation(data.shape[0])
        data = data[rand_indexes]
        label = label[rand_indexes]

    batch_count = len(data) / batch_size
    batches_data = np.split(data[:batch_count * batch_size], batch_count)
    batches_data.append(data[batch_count * batch_size:])
    batches_labels = np.split(label[:batch_count * batch_size], batch_count)
    batches_labels.append(label[batch_count * batch_size:])
    batch_count += 1

    return batches_data, batches_labels, batch_count

def count_params():
    total_params = 0
    for variable in tf.trainable_variables():
        shape = variable.get_shape()
        params = 1
        for dim in shape:
            params = params * dim.value
        total_params += params
    print("Total training params: %.2fM" % (total_params / 1e6))

if __name__ == '__main__':
    ##
    train_params = {'normalize_type': 'by-channels',  # by-channels, divide-255, divide-256
                    'initial_lr': 0.1,
                    'weight_decay': 1e-4,
                    'batch_size': 64,
                    'total_epoch': 300,
                    'keep_prob': 0.8
                    }

    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu', default="0")
    parser.add_argument('--dataset', default='cifar-10', choices=['cifar-10', 'cifar-100', 'svhn'])
    parser.add_argument('--k', default=12, type=int, help='filters per layer')
    parser.add_argument('--T', default=5,type=int, help='total layers in all blocks')
    parser.add_argument('--dir', default='checkpoint/', help='folder to store models')
    parser.add_argument('--if_a', default=False, type=bool,
                        help='if use attentional transition')
    parser.add_argument('--if_b', default=False, type=bool,
                        help='if use bottleneck architecture')
    parser.add_argument('--if_c', default=False, type=bool,
                        help='if use compression')

    args = parser.parse_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    dataset = args.dataset

    if dataset == 'svhn':
        total_epoches = 40
    else:
        total_epoches = train_params['total_epoch']

    result_dir = args.dir

    batch_size = train_params['batch_size']
    lr = train_params['initial_lr']
    kp = train_params['keep_prob']
    weight_decay = train_params['weight_decay']

    if os.path.exists(result_dir) == False:
        os.mkdir(result_dir)

    train_data, train_label, test_data, test_label = load_data(
        dataset, train_params['normalize_type'])

    image_size = train_data.shape[1:]
    label_num = train_label.shape[-1]

    graph = tf.Graph()
    with graph.as_default():
        input_images = tf.placeholder(tf.float32, [
            None, image_size[0], image_size[1], image_size[2]], name='input_images')
        true_labels = tf.placeholder(
            tf.float32, [None, label_num], name='labels')
        is_train = tf.placeholder(tf.bool, shape=[])
        learning_rate = tf.placeholder(
            tf.float32, shape=[], name='learning_rate')
        keep_prob = tf.placeholder(tf.float32, shape=[], name='keep_prob')
        # build model
        logits, prob = build_model(input_images, args.k, args.T, label_num,
                                   is_train, keep_prob, args.if_a, args.if_b, args.if_c)
        ###     loss and accuracy
        loss_cross_entropy = tf.reduce_mean(
            tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=true_labels))
        if_correct = tf.equal(tf.argmax(prob, 1), tf.argmax(true_labels, 1))
        accuracy = tf.reduce_mean(tf.cast(if_correct, tf.float32))
        l2_loss = tf.add_n([tf.nn.l2_loss(var)
                            for var in tf.trainable_variables()])
        # optimizer
        optimizer = tf.train.MomentumOptimizer(
            learning_rate, 0.9, use_nesterov=True)
        train_op = optimizer.minimize(
            loss_cross_entropy + l2_loss * weight_decay)
        saver = tf.train.Saver()

    ###     begin training     ###

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    with tf.Session(config=config, graph=graph) as sess:
        sess.run(tf.global_variables_initializer())
        count_params()
        ###   train batch data   ###

        #  to shuffle before each epoch

        ###   test batch data   ###

        batches_data_test, batches_labels_test, batch_count_test = into_batch(
            test_data, test_label, batch_size, shuffle=False)

        loss_train = []
        acc_train = []
        loss_test = []
        acc_test = []
        best_acc = 0
        for epoch in range(1, total_epoches + 1):
            if epoch == total_epoches / 2:
                lr = lr * 0.1
            if epoch == total_epoches * 3 / 4:
                lr = lr * 0.1

            batches_data, batches_labels, batch_count = into_batch(
                train_data, train_label, batch_size, shuffle=True)

            ###     train     ###
            loss_per_bat = []
            acc_per_bat = []
            for batch_id in range(batch_count):
                data_per_bat = batches_data[batch_id]
                label_per_bat = batches_labels[batch_id]
                result_per_bat = sess.run([train_op, loss_cross_entropy, accuracy],
                                          feed_dict={input_images: data_per_bat,
                                                     true_labels: label_per_bat,
                                                     learning_rate: lr,
                                                     is_train: True,
                                                     keep_prob: kp})
                loss_per_bat.append(result_per_bat[1])
                acc_per_bat.append(result_per_bat[2])
                if (batch_id + 1) % 100 == 0:
                    print('epoch:', epoch, 'batch:',
                          batch_id + 1, 'in', batch_count)
                    print('loss:', result_per_bat[1],
                          'accuracy:', result_per_bat[2])

            saver.save(sess, os.path.join(
                result_dir, dataset + '_epoch_%d.ckpt' % epoch))
            loss_train.append(np.mean(loss_per_bat))
            acc_train.append(np.mean(acc_per_bat))

            ###     test     ###
            loss_per_bat = []
            acc_per_bat = []
            for batch_id in range(batch_count_test):
                data_per_bat = batches_data_test[batch_id]
                label_per_bat = batches_labels_test[batch_id]
                result_per_bat = sess.run([loss_cross_entropy, accuracy],
                                          feed_dict={input_images: data_per_bat,
                                                     true_labels: label_per_bat,
                                                     is_train: False,
                                                     keep_prob: 1})
                loss_per_bat.append(result_per_bat[0])  # result[0]->loss
                acc_per_bat.append(result_per_bat[1])  # result[1]->acc
            loss_test.append(np.mean(loss_per_bat))
            acc_test.append(np.mean(acc_per_bat))

            if acc_test[-1] > best_acc:
                best_acc = acc_test[-1]

            print(time.ctime())
            print('epoch:', epoch)
            print('train loss:', loss_train[-1], 'acc:', acc_train[-1])
            print('test loss:', loss_test[-1], 'acc:', acc_test[-1])
            print('best test acc:', best_acc)
            print('\n')

        np.save(os.path.join(result_dir, result_dir +
                             '_loss_train.npy'), np.array(loss_train))
        np.save(os.path.join(result_dir, result_dir +
                             '_acc_train.npy'), np.array(acc_train))
        np.save(os.path.join(result_dir, result_dir +
                             '_loss_test.npy'), np.array(loss_test))
        np.save(os.path.join(result_dir, result_dir +
                             '_acc_test.npy'), np.array(acc_test))


"""An example to test diayn as task proposal to pearl written in PyTorch."""

import altair as alt
import os
import click
import gym
import joblib
import numpy as np
import pandas as pd
import torch
from altair_saver import save
from torch import nn
from torch.nn import functional as F

import garage.torch.utils as tu
from garage import wrap_experiment
from garage.envs import GarageEnv, DiaynEnvWrapper
from garage.envs import normalize
from garage.envs.mujoco import HalfCheetahVelEnv
from garage.experiment import deterministic, LocalRunner
from garage.experiment.deterministic import set_seed
from garage.experiment.task_sampler import EnvPoolSampler, SetTaskSampler
from garage.replay_buffer import PathBuffer
from garage.sampler import LocalSampler
from garage.sampler import SkillWorker
from garage.sampler.local_skill_sampler import LocalSkillSampler
from garage.torch.algos import DIAYN
from garage.torch.algos import PEARL
from garage.torch.algos.discriminator import MLPDiscriminator
from garage.torch.algos.pearl import PEARLWorker
from garage.torch.embeddings import MLPEncoder
from garage.torch.policies import ContextConditionedPolicy
from garage.torch.policies import TanhGaussianMLPPolicy
from garage.torch.policies import TanhGaussianMLPSkillPolicy
from garage.torch.q_functions import ContinuousMLPQFunction
from garage.torch.q_functions import ContinuousMLPSkillQFunction

skills_num = 10

@wrap_experiment(snapshot_mode='gap_and_last', snapshot_gap=100)
def diayn_half_cheetah_vel_batch_for_pearl(ctxt=None, seed=1):
    deterministic.set_seed(seed)
    runner = LocalRunner(snapshot_config=ctxt)
    env = GarageEnv(normalize(HalfCheetahVelEnv()))

    policy = TanhGaussianMLPSkillPolicy(
        env_spec=env.spec,
        skills_num=skills_num,
        hidden_sizes=[256, 256],
        hidden_nonlinearity=nn.ReLU,
        output_nonlinearity=None,
        min_std=np.exp(-20.),
        max_std=np.exp(2.),
    )

    qf1 = ContinuousMLPSkillQFunction(env_spec=env.spec,
                                      skills_num=skills_num,
                                      hidden_sizes=[256, 256],
                                      hidden_nonlinearity=F.relu)

    qf2 = ContinuousMLPSkillQFunction(env_spec=env.spec,
                                      skills_num=skills_num,
                                      hidden_sizes=[256, 256],
                                      hidden_nonlinearity=F.relu)

    discriminator = MLPDiscriminator(env_spec=env.spec,
                                     skills_num=skills_num,
                                     hidden_sizes=[64, 64],
                                     hidden_nonlinearity=F.relu)

    replay_buffer = PathBuffer(capacity_in_transitions=int(1e6))

    diayn = DIAYN(env_spec=env.spec,
                  skills_num=skills_num,
                  discriminator=discriminator,
                  policy=policy,
                  qf1=qf1,
                  qf2=qf2,
                  gradient_steps_per_itr=1000,
                  max_path_length=300,
                  replay_buffer=replay_buffer,
                  min_buffer_size=1e4,
                  recorded=True,  # enable the video recording func
                  target_update_tau=5e-3,
                  discount=0.99,
                  buffer_batch_size=256,
                  reward_scale=1.,
                  steps_per_epoch=1)

    if torch.cuda.is_available():
        tu.set_gpu_mode(True)
    else:
        tu.set_gpu_mode(False)
    diayn.to()
    worker_args = {"skills_num": skills_num}
    runner.setup(algo=diayn, env=env, sampler_cls=LocalSkillSampler,
                 worker_class=SkillWorker, worker_args=worker_args)
    runner.train(n_epochs=1000, batch_size=1000)  # 1000
    # runner.restore(from_dir=os.path.join(os.getcwd(), 'data/local/experiment/diayn_half_cheetah_batch_50'))
    # diayn = runner.get_algo()
    runner.save(999)  # saves the last episode

    return discriminator, diayn

s = np.random.randint(0, 1000)  # 521 in the sac_cheetah example
task_proposer, diayn_trained_agent = diayn_half_cheetah_vel_batch_for_pearl(seed=s)

'''
load_dir = os.path.join(os.getcwd(), 'data/local/experiment/diayn_half_cheetah_vel_batch_for_pearl_3')
itr = 900
load_from_file = os.path.join(load_dir, 'itr_{}.pkl'.format(itr))
file = open(load_from_file, 'rb')
saved = joblib.load(file)
file.close()
diayn = saved['algo']
task_proposer = diayn.networks[1]  # _discriminator
'''
########################## hyper params for PEARL ##########################

param_num_epoches = 500
param_train_tasks_num = skills_num  # 100
param_test_tasks_num = 5 # skills_num / 2  # 30
param_encoder_hidden_size = 200
param_net_size = 300
param_num_steps_per_epoch = 2000
param_num_initial_steps = 2000
param_num_steps_prior = 400
param_num_extra_rl_steps_posterior = 600
param_batch_size = 256
param_embedding_batch_size = 100
param_embedding_mini_batch_size = 100
param_max_path_length = 300

param_latent_size = 5
param_num_tasks_sample = 5
param_meta_batch_size = 16
param_reward_scale = 5.
param_use_gpu = True

###########################################################################

@click.command()
@click.option('--num_epochs', default=param_num_epoches)
@click.option('--num_train_tasks', default=param_train_tasks_num)
@click.option('--num_test_tasks', default=param_test_tasks_num)
@click.option('--encoder_hidden_size', default=param_encoder_hidden_size)
@click.option('--net_size', default=param_net_size)
@click.option('--num_steps_per_epoch', default=param_num_steps_per_epoch)
@click.option('--num_initial_steps', default=param_num_initial_steps)
@click.option('--num_steps_prior', default=param_num_steps_prior)
@click.option('--num_extra_rl_steps_posterior',
              default=param_num_extra_rl_steps_posterior)
@click.option('--batch_size', default=param_batch_size)
@click.option('--embedding_batch_size', default=param_embedding_batch_size)
@click.option('--embedding_mini_batch_size',
              default=param_embedding_mini_batch_size)
@click.option('--max_path_length', default=param_max_path_length)
@wrap_experiment
def diayn_pearl_half_cheeth(
    ctxt=None,
    seed=1,
    num_epochs=param_num_epoches,
    num_train_tasks=param_train_tasks_num,
    num_test_tasks=param_test_tasks_num,
    latent_size=param_latent_size,
    encoder_hidden_size=param_encoder_hidden_size,
    net_size=param_net_size,
    meta_batch_size=param_meta_batch_size,
    num_steps_per_epoch=param_num_steps_per_epoch,
    num_initial_steps=param_num_initial_steps,
    num_tasks_sample=param_num_tasks_sample,
    num_steps_prior=param_num_steps_prior,
    num_extra_rl_steps_posterior=param_num_extra_rl_steps_posterior,
    batch_size=param_batch_size,
    embedding_batch_size=param_embedding_batch_size,
    embedding_mini_batch_size=param_embedding_mini_batch_size,
    max_path_length=param_max_path_length,
    reward_scale=param_reward_scale,
    use_gpu=param_use_gpu):
    if task_proposer is None:
        raise ValueError("Task proposer is empty")

    assert num_train_tasks is skills_num

    set_seed(seed)
    encoder_hidden_sizes = (encoder_hidden_size, encoder_hidden_size,
                            encoder_hidden_size)
    # create multi-task environment and sample tasks

    ML_train_envs = [DiaynEnvWrapper(task_proposer, skills_num, task_name,
                                     normalize(HalfCheetahVelEnv()))
                     for task_name in range(skills_num)]
    env_sampler = EnvPoolSampler(ML_train_envs)
    env = env_sampler.sample(num_train_tasks)

    # train_trajs_dist = [train_env.get_training_traj(diayn_trained_agent)
    #               for train_env in ML_train_envs]

    # ML_test_envs = [
    #     GarageEnv(normalize(
    #         DiaynEnvWrapper(env, task_proposer, skills_num, task_name)))
    #     for task_name in random.sample(range(skills_num), test_tasks_num)
    # ]

    test_env_sampler = SetTaskSampler(lambda: GarageEnv(normalize(
        HalfCheetahVelEnv())))

    runner = LocalRunner(ctxt)

    # instantiate networks
    augmented_env = PEARL.augment_env_spec(env[0](), latent_size)
    qf = ContinuousMLPQFunction(env_spec=augmented_env,
                                hidden_sizes=[net_size, net_size, net_size])

    vf_env = PEARL.get_env_spec(env[0](), latent_size, 'vf')
    vf = ContinuousMLPQFunction(env_spec=vf_env,
                                hidden_sizes=[net_size, net_size, net_size])

    inner_policy = TanhGaussianMLPPolicy(
        env_spec=augmented_env, hidden_sizes=[net_size, net_size, net_size])

    pearl = PEARL(
        env=env,
        policy_class=ContextConditionedPolicy,
        encoder_class=MLPEncoder,
        inner_policy=inner_policy,
        qf=qf,
        vf=vf,
        num_train_tasks=num_train_tasks,
        num_test_tasks=num_test_tasks,
        latent_dim=latent_size,
        encoder_hidden_sizes=encoder_hidden_sizes,
        test_env_sampler=test_env_sampler,
        meta_batch_size=meta_batch_size,
        num_steps_per_epoch=num_steps_per_epoch,
        num_initial_steps=num_initial_steps,
        num_tasks_sample=num_tasks_sample,
        num_steps_prior=num_steps_prior,
        num_extra_rl_steps_posterior=num_extra_rl_steps_posterior,
        batch_size=batch_size,
        embedding_batch_size=embedding_batch_size,
        embedding_mini_batch_size=embedding_mini_batch_size,
        max_path_length=max_path_length,
        reward_scale=reward_scale,
    )

    tu.set_gpu_mode(use_gpu, gpu_id=0)
    if use_gpu:
        pearl.to()

    runner.setup(algo=pearl,
                 env=env[0](),
                 sampler_cls=LocalSampler,
                 sampler_args=dict(max_path_length=max_path_length),
                 n_workers=1,
                 worker_class=PEARLWorker)

    average_returns = runner.train(n_epochs=num_epochs, batch_size=batch_size)
    runner.save(num_epochs-1)

    return average_returns

@click.command()
@click.option('--num_epochs', default=param_num_epoches)
@click.option('--num_train_tasks', default=param_train_tasks_num)
@click.option('--num_test_tasks', default=param_test_tasks_num)
@click.option('--encoder_hidden_size', default=param_encoder_hidden_size)
@click.option('--net_size', default=param_net_size)
@click.option('--num_steps_per_epoch', default=param_num_steps_per_epoch)
@click.option('--num_initial_steps', default=param_num_initial_steps)
@click.option('--num_steps_prior', default=param_num_steps_prior)
@click.option('--num_extra_rl_steps_posterior',
              default=param_num_extra_rl_steps_posterior)
@click.option('--batch_size', default=param_batch_size)
@click.option('--embedding_batch_size',
              default=param_embedding_mini_batch_size)
@click.option('--embedding_mini_batch_size',
              default=param_embedding_mini_batch_size)
@click.option('--max_path_length', default=param_max_path_length)
@wrap_experiment
def pearl_half_cheetah(ctxt=None,
                       seed=1,
                       num_epochs=param_num_epoches,
                       num_train_tasks=param_train_tasks_num,
                       num_test_tasks=param_test_tasks_num,
                       latent_size=param_latent_size,
                       encoder_hidden_size=param_encoder_hidden_size,
                       net_size=param_net_size,
                       meta_batch_size=param_meta_batch_size,
                       num_steps_per_epoch=param_num_steps_per_epoch,
                       num_initial_steps=param_num_initial_steps,
                       num_tasks_sample=param_num_tasks_sample,
                       num_steps_prior=param_num_steps_prior,
                       num_extra_rl_steps_posterior=param_num_extra_rl_steps_posterior,
                       batch_size=param_batch_size,
                       embedding_batch_size=param_embedding_batch_size,
                       embedding_mini_batch_size=param_embedding_mini_batch_size,
                       max_path_length=param_max_path_length,
                       reward_scale=param_reward_scale,
                       use_gpu=param_use_gpu):
    set_seed(seed)
    encoder_hidden_sizes = (encoder_hidden_size, encoder_hidden_size,
                            encoder_hidden_size)
    # create multi-task environment and sample tasks
    env_sampler = SetTaskSampler(lambda: GarageEnv(
        normalize(HalfCheetahVelEnv())))
    env = env_sampler.sample(num_train_tasks)
    test_env_sampler = SetTaskSampler(lambda: GarageEnv(
        normalize(HalfCheetahVelEnv())))

    runner = LocalRunner(ctxt)

    # instantiate networks
    augmented_env = PEARL.augment_env_spec(env[0](), latent_size)
    qf = ContinuousMLPQFunction(env_spec=augmented_env,
                                hidden_sizes=[net_size, net_size, net_size])

    vf_env = PEARL.get_env_spec(env[0](), latent_size, 'vf')
    vf = ContinuousMLPQFunction(env_spec=vf_env,
                                hidden_sizes=[net_size, net_size, net_size])

    inner_policy = TanhGaussianMLPPolicy(
        env_spec=augmented_env, hidden_sizes=[net_size, net_size, net_size])

    pearl = PEARL(
        env=env,
        policy_class=ContextConditionedPolicy,
        encoder_class=MLPEncoder,
        inner_policy=inner_policy,
        qf=qf,
        vf=vf,
        num_train_tasks=num_train_tasks,
        num_test_tasks=num_test_tasks,
        latent_dim=latent_size,
        encoder_hidden_sizes=encoder_hidden_sizes,
        test_env_sampler=test_env_sampler,
        meta_batch_size=meta_batch_size,
        num_steps_per_epoch=num_steps_per_epoch,
        num_initial_steps=num_initial_steps,
        num_tasks_sample=num_tasks_sample,
        num_steps_prior=num_steps_prior,
        num_extra_rl_steps_posterior=num_extra_rl_steps_posterior,
        batch_size=batch_size,
        embedding_batch_size=embedding_batch_size,
        embedding_mini_batch_size=embedding_mini_batch_size,
        max_path_length=max_path_length,
        reward_scale=reward_scale,
    )

    tu.set_gpu_mode(use_gpu, gpu_id=0)
    if use_gpu:
        pearl.to()

    runner.setup(algo=pearl,
                 env=env[0](),
                 sampler_cls=LocalSampler,
                 sampler_args=dict(max_path_length=max_path_length),
                 n_workers=1,
                 worker_class=PEARLWorker)

    average_returns = runner.train(n_epochs=num_epochs, batch_size=batch_size)
    runner.save(num_epochs - 1)

    return average_returns

def save_list_to_file(x, filename):
    with open(filename, 'w') as f:
        for item in x:
            f.write("%s\n" % item)

if not os.path.exists('tmp'):
    os.makedirs('tmp')
diayn_pearl_returns = diayn_pearl_half_cheeth()
save_list_to_file(diayn_pearl_returns, "tmp/diayn_pearl_returns.txt")
"""
pearl_returns = pearl_half_cheetah()
save_list_to_file(pearl_returns, "tmp/pearl_returns.txt")

assert (len(diayn_pearl_returns) == len(pearl_returns))

n_subject = 2
save_dir = "tmp/avg_rewards.png"
data = pd.DataFrame(
    {'Epoch': [i for i in range(len(diayn_pearl_returns) * 3)] * n_subject,
     'Subjects': ["PEARL with DIAYN" for _ in range(len(diayn_pearl_returns))]
                 + ["PEARL" for _ in range(len(pearl_returns))],
     'Average Rewards': diayn_pearl_returns + pearl_returns})
chart = alt.Chart(data).mark_line().encode(
    x='Epoch:Q',
    y='Average Rewards:Q',
    color='Subjects:N')
save(chart, save_dir)
"""


# -*- coding: utf-8 -*-

'''
@author: hzw77, gjz5038

Network agent from agent.py

'''

import numpy as np
from keras.layers import Input, Dense, Conv2D, Flatten, BatchNormalization, Activation, Multiply, Add
from keras.models import Model, model_from_json, load_model
from keras.optimizers import RMSprop
from keras.layers.core import Dropout
from keras.layers.pooling import MaxPooling2D
from keras import backend as K
import random
from keras.engine.topology import Layer
import os

from agent import Agent, State

class Selector(Layer):

    def __init__(self, select, **kwargs):
        super(Selector, self).__init__(**kwargs)
        self.select = select
        self.select_neuron = K.constant(value=self.select)

    def build(self, input_shape):
        # Create a trainable weight variable for this layer.
        super(Selector, self).build(input_shape)  # Be sure to call this somewhere!

    def call(self, x):
        return K.cast(K.equal(x, self.select_neuron), dtype="float32")

    def get_config(self):
        config = {"select": self.select}
        base_config = super(Selector, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def compute_output_shape(self, input_shape):
        return input_shape

def conv2d_bn(input_layer, index_layer,
              filters=16,
              kernel_size=(3, 3),
              strides=(1, 1)):
    """Utility function to apply conv + BN.
    # Arguments
        x: input tensor.
        filters: filters in `Conv2D`.
        num_row: height of the convolution kernel.
        num_col: width of the convolution kernel.
        padding: padding mode in `Conv2D`.
        strides: strides in `Conv2D`.
        name: name of the ops; will become `name + '_conv'`
            for the convolution and `name + '_bn'` for the
            batch norm layer.
    # Returns
        Output tensor after applying `Conv2D` and `BatchNormalization`.
    """
    if K.image_data_format() == 'channels_first':
        bn_axis = 1
    else:
        bn_axis = 3
    conv = Conv2D(filters=filters,
                  kernel_size=kernel_size,
                  strides=strides,
                  padding='same',
                  use_bias=False,
                  name="conv{0}".format(index_layer))(input_layer)
    bn = BatchNormalization(axis=bn_axis, scale=False, name="bn{0}".format(index_layer))(conv)
    act = Activation('relu', name="act{0}".format(index_layer))(bn)
    pooling = MaxPooling2D(pool_size=2)(act)
    x = Dropout(0.3)(pooling)
    return x

class NetworkAgent(Agent):

    @staticmethod
    def _unison_shuffled_copies(Xs, Y, sample_weight):
        p = np.random.permutation(len(Y))
        new_Xs = []
        for x in Xs:
            assert len(x) == len(Y)
            new_Xs.append(x[p])
        return new_Xs, Y[p], sample_weight[p]

    @staticmethod
    def _cnn_network_structure(img_features):
        conv1 = conv2d_bn(img_features, 1, filters=32, kernel_size=(8, 8), strides=(4, 4))
        conv2 = conv2d_bn(conv1, 2, filters=16, kernel_size=(4, 4), strides=(2, 2))
        img_flatten = Flatten()(conv2)
        return img_flatten

    @staticmethod
    def _shared_network_structure(state_features, dense_d):
        hidden_1 = Dense(dense_d, activation="sigmoid", name="hidden_shared_1")(state_features)
        return hidden_1

    @staticmethod
    def _separate_network_structure(state_features, dense_d, num_actions, memo=""):
        hidden_1 = Dense(dense_d, activation="sigmoid", name="hidden_separate_branch_{0}_1".format(memo))(state_features)
        q_values = Dense(num_actions, activation="linear", name="q_values_separate_branch_{0}".format(memo))(hidden_1)
        return q_values

    def load_model(self, file_name):
        self.q_network = load_model(os.path.join(self.path_set.PATH_TO_MODEL, "%s_q_network.h5" % file_name))

    def save_model(self, file_name):
        self.q_network.save(os.path.join(self.path_set.PATH_TO_MODEL, "%s_q_network.h5" % file_name))

    def choose(self, count, if_pretrain):

        ''' choose the best action for current state '''

        q_values = self.q_network.predict(self.convert_state_to_input(self.state))
        #print(q_values)
        if if_pretrain:
            self.action = np.argmax(q_values[0])
        else:
            if random.random() <= self.para_set.EPSILON:  # continue explore new Random Action
                self.action = random.randrange(len(q_values[0]))
                print("##Explore")
            else:  # exploitation
                self.action = np.argmax(q_values[0])
            if self.para_set.EPSILON > 0.001 and count >= 20000:
                self.para_set.EPSILON = self.para_set.EPSILON * 0.9999
        return self.action, q_values

    def build_memory(self):

        return []

    def build_network_from_copy(self, network_copy):

        '''Initialize a Q network from a copy'''

        network_structure = network_copy.to_json()
        network_weights = network_copy.get_weights()
        network = model_from_json(network_structure, custom_objects={"Selector": Selector})
        network.set_weights(network_weights)
        network.compile(optimizer=RMSprop(lr=self.para_set.LEARNING_RATE),
                        loss="mean_squared_error")
        return network

    #
    # def batch_predict(self,file_name="temp"):
    #     f_samples = open("./records/DQN_v1/" + file_name[:file_name.rfind("_")] + "predict_pretrain.txt", "a")
    #     f_samples_head = ["state.cur_phase", "state.time_this_phase",
    #                       "target",
    #                       "action",
    #                       "reward"]
    #     f_samples.write('\t'.join(f_samples_head) + "\n")
    #     len_memory = len(self.memory)
    #     for i in range(len_memory):
    #         state, action, reward, next_state = self.memory[i]
    #         q_values = self.q_network.predict(self.convert_state_to_input(state))
    #         f_samples.write("{0}\t{1}\t{2}\t{3}\t{4}\n".format(
    #             str(state.cur_phase[0]), str(state.time_this_phase), str(q_values),
    #             str(action), str(reward)
    #         ))

    def remember(self, state, action, reward, next_state):

        ''' log the history '''
        self.memory.append([state, action, reward, next_state])

    def forget(self):

        ''' remove the old history if the memory is too large '''

        if len(self.memory) > self.para_set.MAX_MEMORY_LEN:
            print("length of memory: {0}, before forget".format(len(self.memory)))
            self.memory = self.memory[-self.para_set.MAX_MEMORY_LEN:]
            print("length of memory: {0}, after forget".format(len(self.memory)))

    def _get_next_estimated_reward(self, next_state):

        if self.para_set.DDQN:
            a_max = np.argmax(self.q_network.predict(
                self.convert_state_to_input(next_state))[0])
            next_estimated_reward = self.q_network_bar.predict(
                self.convert_state_to_input(next_state))[0][a_max]
            return next_estimated_reward
        else:
            next_estimated_reward = np.max(self.q_network_bar.predict(
                self.convert_state_to_input(next_state))[0])
            return next_estimated_reward

    def update_network_bar(self):

        ''' update Q bar '''

        if self.q_bar_outdated >= self.para_set.UPDATE_Q_BAR_FREQ:
            self.q_network_bar = self.build_network_from_copy(self.q_network)
            self.q_bar_outdated = 0


from typing import Any

import albumentations as A
from albumentations.pytorch import ToTensorV2
import cv2
import numpy as np
from torch import Tensor

__all__ = ["DINOAugmentation"]

class DINOAugmentation(A.ImageOnlyTransform):
    def __init__(
        self,
        global_crops_scale: tuple[float, float],
        local_crops_scale: tuple[float, float],
        local_crops_number: int,
    ) -> None:
        super().__init__()
        flip_and_color_jitter = A.Compose(
            [
                A.HorizontalFlip(p=0.5),
                A.ColorJitter(p=0.8, brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1),
                A.ToGray(p=0.2),
            ]
        )
        normalize = A.Compose(
            [
                # The default per-mean/std values for albumentations.Normalize
                # match the ImageNet ones prescribed by DINO
                A.Normalize(),
                A.ToFloat(),
                ToTensorV2(),
            ]
        )
        self.global_transfo1 = A.Compose(
            [
                A.RandomResizedCrop(
                    height=224, width=224, scale=global_crops_scale, interpolation=cv2.INTER_CUBIC  # type: ignore
                ),
                flip_and_color_jitter,
                A.GaussianBlur(p=1.0, sigma_limit=(0.1, 2)),
                normalize,
            ]
        )
        self.global_transfo2 = A.Compose(
            [
                A.RandomResizedCrop(
                    height=224, width=224, scale=global_crops_scale, interpolation=cv2.INTER_CUBIC  # type: ignore
                ),
                flip_and_color_jitter,
                A.GaussianBlur(p=1.0, sigma_limit=(0.1, 2)),
                A.Solarize(p=0.2),
                normalize,
            ]
        )
        # transformation for the local small crops
        self.local_crops_number = local_crops_number
        self.local_transfo = A.Compose(
            [
                A.RandomResizedCrop(
                    height=96, width=96, scale=local_crops_scale, interpolation=cv2.INTER_CUBIC  # type: ignore
                ),
                flip_and_color_jitter,
                A.GaussianBlur(p=1.0, sigma_limit=(0.1, 2)),
                *normalize,
            ]
        )

    def __call__(self, image: np.ndarray, **kwargs: Any) -> dict[str, list[Tensor]]:
        crops = [
            self.global_transfo1(image=image)["image"],
            self.global_transfo2(image=image)["image"],
        ]
        for _ in range(self.local_crops_number):
            crops.append(self.local_transfo(image=image)["image"])
        return {"image": crops}