
import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch

def  __snake_case							(	__magic_name__						):  # picklable for multiprocessing







				'''simple docstring'''


				return x.sum()

def  __snake_case							(	__magic_name__						):  # picklable for multiprocessing







				'''simple docstring'''


				return i + 1





@dataclass
class   UpperCamelCase_      :







						'''simple docstring'''



						UpperCamelCase      :  int
						UpperCamelCase      :  str





class   UpperCamelCase_      (      __a		):







						'''simple docstring'''



						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->int:
										lowercase         =      {}
										lowercase         =      []
										lowercase         =      1
										lowercase         =      [1, 2]
										lowercase         =      {"a": 1, "b": 2}
										lowercase         =      {"a": [1, 2], "b": [3, 4]}
										lowercase         =      {"a": {"1": 1}, "b": 2}
										lowercase         =      {"a": 1, "b": 2, "c": 3, "d": 4}
										lowercase         =      {}
										lowercase         =      []
										lowercase         =      2
										lowercase         =      [2, 3]
										lowercase         =      {"a": 2, "b": 3}
										lowercase         =      {"a": [2, 3], "b": [4, 5]}
										lowercase         =      {"a": {"1": 2}, "b": 3}
										lowercase         =      {"a": 2, "b": 3, "c": 4, "d": 5}
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__     )		,    lowerCAmelCase__     )

										lowercase         =      2
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    num_proc=lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    num_proc=lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    num_proc=lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    num_proc=lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    num_proc=lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    num_proc=lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    num_proc=lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    num_proc=lowerCAmelCase__     )		,    lowerCAmelCase__     )

										lowercase         =      {"a": np.eye(2     ), "b": np.zeros(3     ), "c": np.ones(2     )}
										lowercase         =      {"a": 2, "b": 0, "c": 2}
										lowercase         =      {
										    "a": np.eye(2     ).astype(lowerCAmelCase__     ),
										    "b": np.zeros(3     ).astype(lowerCAmelCase__     ),
										    "c": np.ones(2     ).astype(lowerCAmelCase__     ),
										}
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    map_numpy=lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(
										    {k: v.tolist() for k, v in map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    map_numpy=lowerCAmelCase__     ).items()}		,    {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}		,    )
										self.assertEqual(map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    map_numpy=lowerCAmelCase__		,    num_proc=lowerCAmelCase__     )		,    lowerCAmelCase__     )
										self.assertEqual(
										    {k: v.tolist() for k, v in map_nested(lowerCAmelCase__		,    lowerCAmelCase__		,    map_numpy=lowerCAmelCase__		,    num_proc=lowerCAmelCase__     ).items()}		,    {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}		,    )
										with self.assertRaises(lowerCAmelCase__     ):  # can't pickle a local lambda
														map_nested(lambda lowerCAmelCase__     : x + 1		,    lowerCAmelCase__		,    num_proc=lowerCAmelCase__     )





						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->Any:
										lowercase         =      {"a": 1, "b": 2}
										lowercase         =      {"a": 3, "b": 4}
										lowercase         =      {"a": 5, "b": 6}
										lowercase         =      sorted([("a", (1, 3, 5)), ("b", (2, 4, 6))]     )
										self.assertEqual(sorted(zip_dict(lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__     )     )		,    lowerCAmelCase__     )





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->Tuple:





										class   UpperCamelCase_      :







																'''simple docstring'''



																UpperCamelCase      :  Optional[int]				   =					'''bar'''




										lowercase         =      Foo()
										self.assertEqual(foo.my_attr		,    "bar"     )
										with temporary_assignment(lowerCAmelCase__		,    "my_attr"		,    "BAR"     ):
														self.assertEqual(foo.my_attr		,    "BAR"     )
										self.assertEqual(foo.my_attr		,    "bar"     )




@pytest.mark.parametrize(
    "iterable_length, num_proc, expected_num_proc"		,    [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]		,    )
def  __snake_case							(	__magic_name__		,    __magic_name__		,    __magic_name__						):







				'''simple docstring'''


				with patch("datasets.utils.py_utils._single_map_nested"						) as mock_single_map_nested, patch(
				    "datasets.parallel.parallel.Pool"						) as mock_multiprocessing_pool:
								lowercase         =      {F'''{i}''': i for i in range(__magic_name__						)}
								lowercase         =      map_nested(lambda __magic_name__						: x + 10		,    __magic_name__		,    num_proc=__magic_name__		,    parallel_min_length=16						)
								if expected_num_proc == 1:
												assert mock_single_map_nested.called
												assert not mock_multiprocessing_pool.called
								else:
												assert not mock_single_map_nested.called
												assert mock_multiprocessing_pool.called
												assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc





class   UpperCamelCase_      (      __a		):







						'''simple docstring'''



						@require_tf
						def        SCREAMING_SNAKE_CASE(			self							:List[Any]     )  ->Union[str, Any]:
										import tensorflow as tf
										from tensorflow.keras import layers

										lowercase         =      layers.Dense(2     )

										def gen_random_output():
														lowercase         =      tf.random.uniform((1, 3)     )
														return model(lowerCAmelCase__     ).numpy()

										with temp_seed(42		,    set_tensorflow=lowerCAmelCase__     ):
														lowercase         =      gen_random_output()
										with temp_seed(42		,    set_tensorflow=lowerCAmelCase__     ):
														lowercase         =      gen_random_output()
										lowercase         =      gen_random_output()

										np.testing.assert_equal(lowerCAmelCase__		,    lowerCAmelCase__     )
										self.assertGreater(np.abs(outa - outa     ).sum()		,    0     )





						@require_torch
						def        SCREAMING_SNAKE_CASE(			self							:Dict     )  ->Any:
										import torch

										def gen_random_output():
														lowercase         =      torch.nn.Linear(3		,    2     )
														lowercase         =      torch.rand(1		,    3     )
														return model(lowerCAmelCase__     ).detach().numpy()

										with temp_seed(42		,    set_pytorch=lowerCAmelCase__     ):
														lowercase         =      gen_random_output()
										with temp_seed(42		,    set_pytorch=lowerCAmelCase__     ):
														lowercase         =      gen_random_output()
										lowercase         =      gen_random_output()

										np.testing.assert_equal(lowerCAmelCase__		,    lowerCAmelCase__     )
										self.assertGreater(np.abs(outa - outa     ).sum()		,    0     )





						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->List[str]:
										def gen_random_output():
														return np.random.rand(1		,    3     )

										with temp_seed(42     ):
														lowercase         =      gen_random_output()
										with temp_seed(42     ):
														lowercase         =      gen_random_output()
										lowercase         =      gen_random_output()

										np.testing.assert_equal(lowerCAmelCase__		,    lowerCAmelCase__     )
										self.assertGreater(np.abs(outa - outa     ).sum()		,    0     )




@pytest.mark.parametrize("input_data"		,    [{}]						)
def  __snake_case							(	__magic_name__						):







				'''simple docstring'''


				lowercase         =      NestedDataStructure(__magic_name__						).data
				assert output_data == input_data

@pytest.mark.parametrize(
    "data, expected_output"		,    [
        ({}, []),
        ([], []),
        ("foo", ["foo"]),
        (["foo", "bar"], ["foo", "bar"]),
        ([["foo", "bar"]], ["foo", "bar"]),
        ([[["foo"], ["bar"]]], ["foo", "bar"]),
        ([[["foo"], "bar"]], ["foo", "bar"]),
        ({"a": 1, "b": 2}, [1, 2]),
        ({"a": [1, 2], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, [4]]}, [1, 2, 3, 4]),
        ({"a": {"1": 1}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": [2]}, [1, 2]),
    ]		,    )
def  __snake_case							(	__magic_name__		,    __magic_name__						):







				'''simple docstring'''


				lowercase         =      NestedDataStructure(__magic_name__						).flatten()
				assert output == expected_output

def  __snake_case							(	):







				'''simple docstring'''


				lowercase         =      A(x=1		,    y="foobar"						)
				lowercase         =      {"x": 1, "y": "foobar"}
				assert asdict(__magic_name__						) == expected_output

				lowercase         =      {"a": {"b": A(x=10		,    y="foo"						)}, "c": [A(x=20		,    y="bar"						)]}
				lowercase         =      {"a": {"b": {"x": 10, "y": "foo"}}, "c": [{"x": 20, "y": "bar"}]}
				assert asdict(__magic_name__						) == expected_output

				with pytest.raises(__magic_name__						):
								asdict([1, A(x=10		,    y="foo"						)]						)

def  __snake_case							(	__magic_name__						):







				'''simple docstring'''


				return text.split()

def  __snake_case							(	__magic_name__						):







				'''simple docstring'''


				yield (time.time(), content)
				time.sleep(2						)
				yield (time.time(), content)




def  __snake_case							(	):







				'''simple docstring'''


				with Pool(2						) as pool:
								lowercase         =      list(iflatmap_unordered(__magic_name__		,    _split_text		,    kwargs_iterable=[{"text": "hello there"}] * 10						)						)
								assert out.count("hello"						) == 10
								assert out.count("there"						) == 10
								assert len(__magic_name__						) == 20

				# check multiprocess from pathos (uses dill for pickling)
				with multiprocess.Pool(2						) as pool:
								lowercase         =      list(iflatmap_unordered(__magic_name__		,    _split_text		,    kwargs_iterable=[{"text": "hello there"}] * 10						)						)
								assert out.count("hello"						) == 10
								assert out.count("there"						) == 10
								assert len(__magic_name__						) == 20

				# check that we get items as fast as possible
				with Pool(2						) as pool:
								lowercase         =      []
								for yield_time, content in iflatmap_unordered(
								    __magic_name__		,    _aseconds_generator_of_aitems_with_timing		,    kwargs_iterable=[{"content": "a"}, {"content": "b"}]						):
												assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
												out.append(__magic_name__						)
								assert out.count("a"						) == 2
								assert out.count("b"						) == 2
								assert len(__magic_name__						) == 4


from typing import List, Optional, Union

import numpy as np

from ....audio_utils import mel_filter_bank, optimal_fft_length, spectrogram, window_function
from ....feature_extraction_sequence_utils import SequenceFeatureExtractor
from ....feature_extraction_utils import BatchFeature
from ....file_utils import PaddingStrategy, TensorType
from ....utils import logging


_snake_case      :							Any     	=    logging.get_logger(__name__)





class   UpperCamelCase_      (      __a		):







						'''simple docstring'''



						UpperCamelCase      :  Any				   =					['''input_features''', '''attention_mask''']





						def __init__(			self							:List[Any]		,    lowerCAmelCase__							:List[Any]=80		,    lowerCAmelCase__							:int=16000		,    lowerCAmelCase__							:str=0.0		,    lowerCAmelCase__							:int=10		,    lowerCAmelCase__							:List[Any]=25		,    lowerCAmelCase__							:List[str]="hamming_window"		,    lowerCAmelCase__							:int=3_27_68.0		,    lowerCAmelCase__							:str=0.97		,    lowerCAmelCase__							:List[str]=1.0		,    lowerCAmelCase__							:Optional[int]=True		,    lowerCAmelCase__							:Tuple=True		,    lowerCAmelCase__							:str=False		,    **lowerCAmelCase__							:Optional[Any]		,    )  ->Union[str, Any]:
										super().__init__(feature_size=lowerCAmelCase__		,    sampling_rate=lowerCAmelCase__		,    padding_value=lowerCAmelCase__		,    **lowerCAmelCase__     )

										lowercase         =      feature_size
										lowercase         =      sampling_rate
										lowercase         =      padding_value
										lowercase         =      hop_length
										lowercase         =      win_length
										lowercase         =      frame_signal_scale
										lowercase         =      preemphasis_coeff
										lowercase         =      mel_floor
										lowercase         =      normalize_means
										lowercase         =      normalize_vars
										lowercase         =      win_function
										lowercase         =      return_attention_mask

										lowercase         =      win_length * sampling_rate // 1000
										lowercase         =      hop_length * sampling_rate // 1000

										lowercase         =      optimal_fft_length(self.sample_size     )
										lowercase         =      (self.n_fft // 2) + 1





						def        SCREAMING_SNAKE_CASE(			self							:Union[str, Any]		,    lowerCAmelCase__							:np.array     )  ->np.ndarray:
										if self.win_function == "hamming_window":
														lowercase         =      window_function(window_length=self.sample_size		,    name=self.win_function		,    periodic=lowerCAmelCase__     )
										else:
														lowercase         =      window_function(window_length=self.sample_size		,    name=self.win_function     )

										lowercase         =      mel_filter_bank(
										    num_frequency_bins=self.n_freqs		,    num_mel_filters=self.feature_size		,    min_frequency=0.0		,    max_frequency=self.sampling_rate / 2.0		,    sampling_rate=self.sampling_rate		,    )

										lowercase         =      spectrogram(
										    one_waveform * self.frame_signal_scale		,    window=lowerCAmelCase__		,    frame_length=self.sample_size		,    hop_length=self.sample_stride		,    fft_length=self.n_fft		,    center=lowerCAmelCase__		,    preemphasis=self.preemphasis_coeff		,    mel_filters=lowerCAmelCase__		,    mel_floor=self.mel_floor		,    log_mel="log"		,    )
										return msfc_features.T





						def        SCREAMING_SNAKE_CASE(			self							:Tuple		,    lowerCAmelCase__							:Optional[Any]		,    lowerCAmelCase__							:Optional[Any]		,    lowerCAmelCase__							:Dict     )  ->int:
										# make sure we normalize float32 arrays
										if self.normalize_means:
														lowercase         =      x[:input_length].mean(axis=0     )
														lowercase         =      np.subtract(lowerCAmelCase__		,    lowerCAmelCase__     )
										if self.normalize_vars:
														lowercase         =      x[:input_length].std(axis=0     )
														lowercase         =      np.divide(lowerCAmelCase__		,    lowerCAmelCase__     )

										if input_length < x.shape[0]:
														lowercase         =      padding_value

										# make sure array is in float32
										lowercase         =      x.astype(np.floataa     )

										return x





						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]		,    lowerCAmelCase__							:List[np.ndarray]		,    lowerCAmelCase__							:Optional[np.ndarray] = None     )  ->List[np.ndarray]:
										lowercase         =      attention_mask.sum(-1     ) if attention_mask is not None else [x.shape[0] for x in input_features]
										return [self._normalize_one(lowerCAmelCase__		,    lowerCAmelCase__		,    self.padding_value     ) for x, n in zip(lowerCAmelCase__		,    lowerCAmelCase__     )]





						def __call__(			self							:Any		,    lowerCAmelCase__							:Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]		,    lowerCAmelCase__							:Union[bool, str, PaddingStrategy] = False		,    lowerCAmelCase__							:Optional[int] = None		,    lowerCAmelCase__							:bool = False		,    lowerCAmelCase__							:Optional[int] = None		,    lowerCAmelCase__							:Optional[bool] = None		,    lowerCAmelCase__							:Optional[Union[str, TensorType]] = None		,    lowerCAmelCase__							:Optional[int] = None		,    **lowerCAmelCase__							:Union[str, Any]		,    )  ->BatchFeature:

										if sampling_rate is not None:
														if sampling_rate != self.sampling_rate:
																		raise ValueError(
																		    F'''The model corresponding to this feature extractor: {self} was trained using a sampling rate of'''
																		    F''' {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with'''
																		    F''' {self.sampling_rate} and not {sampling_rate}.'''     )
										else:
														logger.warning(
														    "It is strongly recommended to pass the ``sampling_rate`` argument to this function. "
														    "Failing to do so can result in silent errors that might be hard to debug."     )

										lowercase         =      isinstance(lowerCAmelCase__		,    np.ndarray     ) and len(raw_speech.shape     ) > 1
										if is_batched_numpy and len(raw_speech.shape     ) > 2:
														raise ValueError(F'''Only mono-channel audio is supported for input to {self}'''     )
										lowercase         =      is_batched_numpy or (
										    isinstance(lowerCAmelCase__		,    (list, tuple)     ) and (isinstance(raw_speech[0]		,    (np.ndarray, tuple, list)     ))
										)

										if is_batched:
														lowercase         =      [np.asarray(lowerCAmelCase__		,    dtype=np.floataa     ) for speech in raw_speech]
										elif not is_batched and not isinstance(lowerCAmelCase__		,    np.ndarray     ):
														lowercase         =      np.asarray(lowerCAmelCase__		,    dtype=np.floataa     )
										elif isinstance(lowerCAmelCase__		,    np.ndarray     ) and raw_speech.dtype is np.dtype(np.floataa     ):
														lowercase         =      raw_speech.astype(np.floataa     )

										# always return batch
										if not is_batched:
														lowercase         =      [raw_speech]

										# extract fbank features
										lowercase         =      [self._extract_mfsc_features(lowerCAmelCase__     ) for one_waveform in raw_speech]

										# convert into correct format for padding
										lowercase         =      BatchFeature({"input_features": features}     )

										lowercase         =      self.pad(
										    lowerCAmelCase__		,    padding=lowerCAmelCase__		,    max_length=lowerCAmelCase__		,    truncation=lowerCAmelCase__		,    pad_to_multiple_of=lowerCAmelCase__		,    return_attention_mask=lowerCAmelCase__		,    **lowerCAmelCase__		,    )
										# make sure list is in array format
										lowercase         =      padded_inputs.get("input_features"     )
										if isinstance(input_features[0]		,    lowerCAmelCase__     ):
														lowercase         =      [np.asarray(lowerCAmelCase__		,    dtype=np.floataa     ) for feature in input_features]

										lowercase         =      padded_inputs.get("attention_mask"     )
										if attention_mask is not None:
														lowercase         =      [np.asarray(lowerCAmelCase__		,    dtype=np.intaa     ) for array in attention_mask]

										if self.normalize_means or self.normalize_vars:
														lowercase         =      (
														    np.array(lowerCAmelCase__		,    dtype=np.intaa     )
														    if self._get_padding_strategies(lowerCAmelCase__		,    max_length=lowerCAmelCase__     ) is not PaddingStrategy.DO_NOT_PAD
														    and padding
														    else None
														)
														lowercase         =      self.normalize(
														    padded_inputs["input_features"]		,    attention_mask=lowerCAmelCase__     )

										if return_tensors is not None:
														lowercase         =      padded_inputs.convert_to_tensors(lowerCAmelCase__     )

										return padded_inputs
