import copy
import os
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import pytest

from datasets.arrow_writer import ArrowWriter, OptimizedTypedSequence, ParquetWriter, TypedSequence
from datasets.features import ArrayaD, ClassLabel, Features, Image, Value
from datasets.features.features import ArrayaDExtensionType, cast_to_python_objects
from datasets.keyhash import DuplicatedKeysError, InvalidKeyError

from .utils import require_pil





class UpperCAmelCase						( __A  ):


				'''simple docstring'''






				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Any         =				pa.array(TypedSequence([1, 2, 3] ) )
						self.assertEqual(arr.type		,     pa.intaa() )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						with self.assertRaises(lowercase ):
								A_		:    Optional[int]         =				pa.array(TypedSequence([1, 2, 3] )		,     type=pa.intaa() )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						with self.assertRaises(lowercase ):
								A_		:    str         =				pa.array(TypedSequence([1, 2, 3]		,     try_type=Value('bool' )		,     type=Value('int64' ) ) )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Dict         =				pa.array(TypedSequence([1, 2, 3]		,     type=Value('int32' ) ) )
						self.assertEqual(arr.type		,     pa.intaa() )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						with self.assertRaises((TypeError, pa.lib.ArrowInvalid) ):
								A_		:    Dict         =				pa.array(TypedSequence(['foo', 'bar']		,     type=Value('int64' ) ) )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Union[str, Any]         =				pa.array(TypedSequence([1, 2, 3]		,     try_type=Value('int32' ) ) )
						self.assertEqual(arr.type		,     pa.intaa() )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Union[str, Any]         =				pa.array(TypedSequence(['foo', 'bar']		,     try_type=Value('int64' ) ) )
						self.assertEqual(arr.type		,     pa.string() )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Optional[int]         =				pa.array(TypedSequence([[[1, 2, 3]]]		,     type=ArrayaD((1, 3)		,     'int64' ) ) )
						self.assertEqual(arr.type		,     ArrayaDExtensionType((1, 3)		,     'int64' ) )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						with self.assertRaises((TypeError, pa.lib.ArrowInvalid) ):
								A_		:    List[str]         =				pa.array(TypedSequence(['foo', 'bar']		,     type=ArrayaD((1, 3)		,     'int64' ) ) )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    List[str]         =				pa.array(TypedSequence([[[1, 2, 3]]]		,     try_type=ArrayaD((1, 3)		,     'int64' ) ) )
						self.assertEqual(arr.type		,     ArrayaDExtensionType((1, 3)		,     'int64' ) )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Optional[Any]         =				pa.array(TypedSequence(['foo', 'bar']		,     try_type=ArrayaD((1, 3)		,     'int64' ) ) )
						self.assertEqual(arr.type		,     pa.string() )







				@require_pil
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						import PIL.Image

						A_		:    List[str]         =				PIL.Image.fromarray(np.arange(1_0		,     dtype=np.uinta ).reshape(2		,     5 ) )
						with patch(
						    'datasets.arrow_writer.cast_to_python_objects'		,     side_effect=lowercase ) as mock_cast_to_python_objects:
								A_		:    Any         =				pa.array(TypedSequence([{'path': None, 'bytes': b'image_bytes'}, pil_image]		,     type=Image() ) )
								A_		:    Union[str, Any]         =				mock_cast_to_python_objects.call_args_list[-1]
								self.assertIn('optimize_list_casting'		,     lowercase )
								self.assertFalse(kwargs['optimize_list_casting'] )




def     UpperCamelCase						(	__lowercase						:			Dict						,__lowercase						:			int     ):
		'''simple docstring'''




		A_		:    Any         =				pa.BufferReader(__lowercase     ) if isinstance(__lowercase						,pa.Buffer     ) else pa.memory_map(__lowercase     )
		A_		:    Tuple         =				pa.ipc.open_stream(__lowercase     )
		A_		:    pa.Table         =				f.read_all()
		assert len(pa_table.to_batches()     ) == expected_num_chunks
		assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}
		del pa_table




@pytest.mark.parametrize('writer_batch_size'						,[None, 1, 10]     )
@pytest.mark.parametrize(
    'fields'						,[None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]     )
def     UpperCamelCase						(	__lowercase						:			Dict						,__lowercase						:			str     ):
		'''simple docstring'''




		A_		:    List[Any]         =				pa.BufferOutputStream()
		A_		:    List[str]         =				pa.schema(__lowercase     ) if fields else None
		with ArrowWriter(stream=__lowercase						,schema=__lowercase						,writer_batch_size=__lowercase     ) as writer:
				writer.write({'col_1': 'foo', 'col_2': 1}     )
				writer.write({'col_1': 'bar', 'col_2': 2}     )
				A_		:    Tuple         =				writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		if not fields:
				A_		:    Union[str, Any]         =				{'col_1': pa.string(), 'col_2': pa.intaa()}
		assert writer._schema == pa.schema(__lowercase						,metadata=writer._schema.metadata     )
		_check_output(output.getvalue()						,expected_num_chunks=num_examples if writer_batch_size == 1 else 1     )




def     UpperCamelCase						(	):
		'''simple docstring'''




		A_		:    Dict         =				pa.BufferOutputStream()
		A_		:    Any         =				Features({'labels': ClassLabel(names=['neg', 'pos']     )}     )
		with ArrowWriter(stream=__lowercase						,features=__lowercase     ) as writer:
				writer.write({'labels': 0}     )
				writer.write({'labels': 1}     )
				A_		:    Optional[Any]         =				writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		assert writer._schema == features.arrow_schema
		assert writer._schema.metadata == features.arrow_schema.metadata
		A_		:    Optional[Any]         =				pa.BufferReader(output.getvalue()     )
		A_		:    List[str]         =				pa.ipc.open_stream(__lowercase     )
		A_		:    pa.Table         =				f.read_all()
		A_		:    List[str]         =				pa_table.schema
		assert pa_table.num_rows == 2
		assert schema == features.arrow_schema
		assert schema.metadata == features.arrow_schema.metadata
		assert features == Features.from_arrow_schema(__lowercase     )




@pytest.mark.parametrize('writer_batch_size'						,[None, 1, 10]     )
def     UpperCamelCase						(	__lowercase						:			List[str]     ):
		'''simple docstring'''




		A_		:    str         =				pa.BufferOutputStream()
		with ArrowWriter(
		    stream=__lowercase						,writer_batch_size=__lowercase						,hash_salt='split_name'						,check_duplicates=__lowercase						,) as writer:
				with pytest.raises(__lowercase     ):
						writer.write({'col_1': 'foo', 'col_2': 1}						,key=[1, 2]     )
						A_		:    Dict         =				writer.finalize()




@pytest.mark.parametrize('writer_batch_size'						,[None, 2, 10]     )
def     UpperCamelCase						(	__lowercase						:			int     ):
		'''simple docstring'''




		A_		:    Dict         =				pa.BufferOutputStream()
		with ArrowWriter(
		    stream=__lowercase						,writer_batch_size=__lowercase						,hash_salt='split_name'						,check_duplicates=__lowercase						,) as writer:
				with pytest.raises(__lowercase     ):
						writer.write({'col_1': 'foo', 'col_2': 1}						,key=10     )
						writer.write({'col_1': 'bar', 'col_2': 2}						,key=10     )
						A_		:    List[Any]         =				writer.finalize()




@pytest.mark.parametrize('writer_batch_size'						,[None, 2, 10]     )
def     UpperCamelCase						(	__lowercase						:			str     ):
		'''simple docstring'''




		A_		:    Optional[int]         =				pa.BufferOutputStream()
		with ArrowWriter(
		    stream=__lowercase						,writer_batch_size=__lowercase						,hash_salt='split_name'						,check_duplicates=__lowercase						,) as writer:
				writer.write({'col_1': 'foo', 'col_2': 1}						,key=1     )
				writer.write({'col_1': 'bar', 'col_2': 2}						,key=2     )
				A_		:    Dict         =				writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		_check_output(output.getvalue()						,expected_num_chunks=num_examples if writer_batch_size == 1 else 1     )




@pytest.mark.parametrize('writer_batch_size'						,[None, 1, 10]     )
@pytest.mark.parametrize(
    'fields'						,[None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]     )
def     UpperCamelCase						(	__lowercase						:			Tuple						,__lowercase						:			str     ):
		'''simple docstring'''




		A_		:    Optional[Any]         =				pa.BufferOutputStream()
		A_		:    Optional[int]         =				pa.schema(__lowercase     ) if fields else None
		with ArrowWriter(stream=__lowercase						,schema=__lowercase						,writer_batch_size=__lowercase     ) as writer:
				writer.write_batch({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}     )
				writer.write_batch({'col_1': [], 'col_2': []}     )
				A_		:    int         =				writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		if not fields:
				A_		:    Optional[Any]         =				{'col_1': pa.string(), 'col_2': pa.intaa()}
		assert writer._schema == pa.schema(__lowercase						,metadata=writer._schema.metadata     )
		_check_output(output.getvalue()						,expected_num_chunks=num_examples if writer_batch_size == 1 else 1     )




@pytest.mark.parametrize('writer_batch_size'						,[None, 1, 10]     )
@pytest.mark.parametrize(
    'fields'						,[None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]     )
def     UpperCamelCase						(	__lowercase						:			int						,__lowercase						:			Optional[int]     ):
		'''simple docstring'''




		A_		:    List[Any]         =				pa.BufferOutputStream()
		A_		:    Tuple         =				pa.schema(__lowercase     ) if fields else None
		with ArrowWriter(stream=__lowercase						,schema=__lowercase						,writer_batch_size=__lowercase     ) as writer:
				writer.write_table(pa.Table.from_pydict({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}     )     )
				A_		:    Optional[int]         =				writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		if not fields:
				A_		:    List[str]         =				{'col_1': pa.string(), 'col_2': pa.intaa()}
		assert writer._schema == pa.schema(__lowercase						,metadata=writer._schema.metadata     )
		_check_output(output.getvalue()						,expected_num_chunks=num_examples if writer_batch_size == 1 else 1     )




@pytest.mark.parametrize('writer_batch_size'						,[None, 1, 10]     )
@pytest.mark.parametrize(
    'fields'						,[None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]     )
def     UpperCamelCase						(	__lowercase						:			Dict						,__lowercase						:			Optional[int]     ):
		'''simple docstring'''




		A_		:    Optional[Any]         =				pa.BufferOutputStream()
		A_		:    Optional[Any]         =				pa.schema(__lowercase     ) if fields else None
		with ArrowWriter(stream=__lowercase						,schema=__lowercase						,writer_batch_size=__lowercase     ) as writer:
				writer.write_row(pa.Table.from_pydict({'col_1': ['foo'], 'col_2': [1]}     )     )
				writer.write_row(pa.Table.from_pydict({'col_1': ['bar'], 'col_2': [2]}     )     )
				A_		:    Tuple         =				writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		if not fields:
				A_		:    Optional[Any]         =				{'col_1': pa.string(), 'col_2': pa.intaa()}
		assert writer._schema == pa.schema(__lowercase						,metadata=writer._schema.metadata     )
		_check_output(output.getvalue()						,expected_num_chunks=num_examples if writer_batch_size == 1 else 1     )




def     UpperCamelCase						(	):
		'''simple docstring'''




		with tempfile.TemporaryDirectory() as tmp_dir:
				A_		:    List[Any]         =				{'col_1': pa.string(), 'col_2': pa.intaa()}
				A_		:    Optional[Any]         =				os.path.join(__lowercase						,'test.arrow'     )
				with ArrowWriter(path=__lowercase						,schema=pa.schema(__lowercase     )     ) as writer:
						writer.write_batch({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}     )
						A_		:    Tuple         =				writer.finalize()
				assert num_examples == 2
				assert num_bytes > 0
				assert writer._schema == pa.schema(__lowercase						,metadata=writer._schema.metadata     )
				_check_output(__lowercase						,1     )




def     UpperCamelCase						(	__lowercase						:			Union[str, Any]     ):
		'''simple docstring'''




		if pa.types.is_list(__lowercase     ):
				return get_base_dtype(arr_type.value_type     )
		else:
				return arr_type




def     UpperCamelCase						(	__lowercase						:			List[Any]						,__lowercase						:			str     ):
		'''simple docstring'''




		if isinstance(lst[0]						,__lowercase     ):
				change_first_primitive_element_in_list(lst[0]						,__lowercase     )
		else:
				A_		:    List[Any]         =				value




@pytest.mark.parametrize('optimized_int_type, expected_dtype'						,[(None, pa.intaa()), (Value('int32'     ), pa.intaa())]     )
@pytest.mark.parametrize('sequence'						,[[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]]     )
def     UpperCamelCase						(	__lowercase						:			Dict						,__lowercase						:			int						,__lowercase						:			Optional[int]     ):
		'''simple docstring'''




		A_		:    Optional[int]         =				pa.array(TypedSequence(__lowercase						,optimized_int_type=__lowercase     )     )
		assert get_base_dtype(arr.type     ) == expected_dtype




@pytest.mark.parametrize(
    'col, expected_dtype'						,[
        ('attention_mask', pa.inta()),
        ('special_tokens_mask', pa.inta()),
        ('token_type_ids', pa.inta()),
        ('input_ids', pa.intaa()),
        ('other', pa.intaa()),
    ]						,)
@pytest.mark.parametrize('sequence'						,[[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]]     )
def     UpperCamelCase						(	__lowercase						:			List[str]						,__lowercase						:			Dict						,__lowercase						:			List[str]     ):
		'''simple docstring'''




		A_		:    Union[str, Any]         =				pa.array(OptimizedTypedSequence(__lowercase						,col=__lowercase     )     )
		assert get_base_dtype(arr.type     ) == expected_dtype

		# not in range
		if col != "other":
				# avoids errors due to in-place modifications
				A_		:    Any         =				copy.deepcopy(__lowercase     )
				A_		:    Union[str, Any]         =				np.iinfo(expected_dtype.to_pandas_dtype()     ).max + 1
				change_first_primitive_element_in_list(__lowercase						,__lowercase     )
				A_		:    Any         =				pa.array(OptimizedTypedSequence(__lowercase						,col=__lowercase     )     )
				assert get_base_dtype(arr.type     ) == pa.intaa()




@pytest.mark.parametrize('raise_exception'						,[False, True]     )
def     UpperCamelCase						(	__lowercase						:			List[str]						,__lowercase						:			Union[str, Any]     ):
		'''simple docstring'''




		A_		:    Dict         =				str(tmp_path / 'dataset-train.arrow'     )
		try:
				with ArrowWriter(path=__lowercase     ) as writer:
						if raise_exception:
								raise pa.lib.ArrowInvalid()
						else:
								writer.stream.close()
		except pa.lib.ArrowInvalid:
				pass
		finally:
				assert writer.stream.closed




def     UpperCamelCase						(	__lowercase						:			List[Any]     ):
		'''simple docstring'''




		A_		:    List[Any]         =				'mock://dataset-train.arrow'
		with ArrowWriter(path=__lowercase						,storage_options=mockfs.storage_options     ) as writer:
				assert isinstance(writer._fs						,type(__lowercase     )     )
				assert writer._fs.storage_options == mockfs.storage_options
				writer.write({'col_1': 'foo', 'col_2': 1}     )
				writer.write({'col_1': 'bar', 'col_2': 2}     )
				A_		:    Optional[int]         =				writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		assert mockfs.exists(__lowercase     )




def     UpperCamelCase						(	):
		'''simple docstring'''




		A_		:    Union[str, Any]         =				pa.BufferOutputStream()
		with ParquetWriter(stream=__lowercase     ) as writer:
				writer.write({'col_1': 'foo', 'col_2': 1}     )
				writer.write({'col_1': 'bar', 'col_2': 2}     )
				A_		:    str         =				writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		A_		:    int         =				pa.BufferReader(output.getvalue()     )
		A_		:    pa.Table         =				pq.read_table(__lowercase     )
		assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}




@require_pil
@pytest.mark.parametrize('embed_local_files'						,[False, True]     )
def     UpperCamelCase						(	__lowercase						:			Optional[Any]						,__lowercase						:			Dict     ):
		'''simple docstring'''




		import PIL.Image

		A_		:    Union[str, Any]         =				str(tmp_path / 'test_image_rgb.jpg'     )
		PIL.Image.fromarray(np.zeros((5, 5)						,dtype=np.uinta     )     ).save(__lowercase						,format='png'     )
		A_		:    List[str]         =				pa.BufferOutputStream()
		with ParquetWriter(
		    stream=__lowercase						,features=Features({'image': Image()}     )						,embed_local_files=__lowercase     ) as writer:
				writer.write({'image': image_path}     )
				writer.finalize()
		A_		:    Tuple         =				pa.BufferReader(output.getvalue()     )
		A_		:    pa.Table         =				pq.read_table(__lowercase     )
		A_		:    Union[str, Any]         =				pa_table.to_pydict()
		if embed_local_files:
				assert isinstance(out['image'][0]['path']						,__lowercase     )
				with open(__lowercase						,'rb'     ) as f:
						assert out["image"][0]["bytes"] == f.read()
		else:
				assert out["image"][0]["path"] == image_path
				assert out["image"][0]["bytes"] is None




def     UpperCamelCase						(	):
		'''simple docstring'''




		A_		:    Union[str, Any]         =				pa.schema([pa.field('col_1'						,pa.string()						,nullable=__lowercase     )]     )
		A_		:    Optional[Any]         =				pa.BufferOutputStream()
		with ArrowWriter(stream=__lowercase     ) as writer:
				writer._build_writer(inferred_schema=__lowercase     )
				assert writer._schema == pa.schema([pa.field('col_1'						,pa.string()     )]     )


from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
	if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.25.0""")):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from ...utils.dummy_torch_and_transformers_objects import UnCLIPImageVariationPipeline, UnCLIPPipeline
else:
	from .pipeline_unclip import UnCLIPPipeline
	from .pipeline_unclip_image_variation import UnCLIPImageVariationPipeline
	from .text_proj import UnCLIPTextProjModel






