
import io
import json

import fsspec
import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.json import JsonDatasetReader, JsonDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases


def 							__A(lowerCAmelCase						,  lowerCAmelCase )		->				Any:

		"""simple docstring"""






		assert isinstance(lowerCAmelCase						,  lowerCAmelCase )
		assert dataset.num_rows == 4
		assert dataset.num_columns == 3
		assert dataset.column_names == ["col_1", "col_2", "col_3"]
		for feature, expected_dtype in expected_features.items():
				assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("""keep_in_memory"""						,  [False, True] )
def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				Union[str, Any]:

		"""simple docstring"""






		_UpperCamelCase						    =	tmp_path / """cache"""
		_UpperCamelCase						    =	{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
		with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
				_UpperCamelCase						    =	JsonDatasetReader(lowerCAmelCase						,  cache_dir=lowerCAmelCase						,  keep_in_memory=lowerCAmelCase ).read()
		_check_json_dataset(lowerCAmelCase						,  lowerCAmelCase )


@pytest.mark.parametrize(
    """features"""						,  [
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]						,  )
def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				Union[str, Any]:

		"""simple docstring"""






		_UpperCamelCase						    =	tmp_path / """cache"""
		_UpperCamelCase						    =	{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
		_UpperCamelCase						    =	features.copy() if features else default_expected_features
		_UpperCamelCase						    =	(
		    Features({feature: Value(lowerCAmelCase ) for feature, dtype in features.items()} ) if features is not None else None
		)
		_UpperCamelCase						    =	JsonDatasetReader(lowerCAmelCase						,  features=lowerCAmelCase						,  cache_dir=lowerCAmelCase ).read()
		_check_json_dataset(lowerCAmelCase						,  lowerCAmelCase )


@pytest.mark.parametrize(
    """features"""						,  [
        None,
        {"""col_3""": """float64""", """col_1""": """string""", """col_2""": """int64"""},
    ]						,  )
def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				str:

		"""simple docstring"""






		_UpperCamelCase						    =	tmp_path / """cache"""
		_UpperCamelCase						    =	{"""col_3""": """float64""", """col_1""": """string""", """col_2""": """int64"""}
		_UpperCamelCase						    =	features.copy() if features else default_expected_features
		_UpperCamelCase						    =	(
		    Features({feature: Value(lowerCAmelCase ) for feature, dtype in features.items()} ) if features is not None else None
		)
		_UpperCamelCase						    =	JsonDatasetReader(lowerCAmelCase						,  features=lowerCAmelCase						,  cache_dir=lowerCAmelCase ).read()
		assert isinstance(lowerCAmelCase						,  lowerCAmelCase )
		assert dataset.num_rows == 2
		assert dataset.num_columns == 3
		assert dataset.column_names == ["col_3", "col_1", "col_2"]
		for feature, expected_dtype in expected_features.items():
				assert dataset.features[feature].dtype == expected_dtype


def 							__A(lowerCAmelCase						,  lowerCAmelCase )		->				Tuple:

		"""simple docstring"""






		_UpperCamelCase						    =	{"""col_2""": """int64""", """col_3""": """float64""", """col_1""": """string"""}
		_UpperCamelCase						    =	features.copy()
		_UpperCamelCase						    =	(
		    Features({feature: Value(lowerCAmelCase ) for feature, dtype in features.items()} ) if features is not None else None
		)
		_UpperCamelCase						    =	tmp_path / """cache"""
		_UpperCamelCase						    =	JsonDatasetReader(lowerCAmelCase						,  features=lowerCAmelCase						,  cache_dir=lowerCAmelCase ).read()
		assert isinstance(lowerCAmelCase						,  lowerCAmelCase )
		assert dataset.num_rows == 2
		assert dataset.num_columns == 3
		assert dataset.column_names == ["col_2", "col_3", "col_1"]
		for feature, expected_dtype in expected_features.items():
				assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("""split"""						,  [None, NamedSplit("""train""" ), """train""", """test"""] )
def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				List[Any]:

		"""simple docstring"""






		_UpperCamelCase						    =	tmp_path / """cache"""
		_UpperCamelCase						    =	{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
		_UpperCamelCase						    =	JsonDatasetReader(lowerCAmelCase						,  cache_dir=lowerCAmelCase						,  split=lowerCAmelCase ).read()
		_check_json_dataset(lowerCAmelCase						,  lowerCAmelCase )
		assert dataset.split == split if split else "train"


@pytest.mark.parametrize("""path_type"""						,  [str, list] )
def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				List[Any]:

		"""simple docstring"""






		if issubclass(lowerCAmelCase						,  lowerCAmelCase ):
				_UpperCamelCase						    =	jsonl_path
		elif issubclass(lowerCAmelCase						,  lowerCAmelCase ):
				_UpperCamelCase						    =	[jsonl_path]
		_UpperCamelCase						    =	tmp_path / """cache"""
		_UpperCamelCase						    =	{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
		_UpperCamelCase						    =	JsonDatasetReader(lowerCAmelCase						,  cache_dir=lowerCAmelCase ).read()
		_check_json_dataset(lowerCAmelCase						,  lowerCAmelCase )


def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase=("train",) )		->				Optional[Any]:

		"""simple docstring"""






		assert isinstance(lowerCAmelCase						,  lowerCAmelCase )
		for split in splits:
				_UpperCamelCase						    =	dataset_dict[split]
				assert dataset.num_rows == 4
				assert dataset.num_columns == 3
				assert dataset.column_names == ["col_1", "col_2", "col_3"]
				for feature, expected_dtype in expected_features.items():
						assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("""keep_in_memory"""						,  [False, True] )
def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				List[Any]:

		"""simple docstring"""






		_UpperCamelCase						    =	tmp_path / """cache"""
		_UpperCamelCase						    =	{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
		with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
				_UpperCamelCase						    =	JsonDatasetReader({"""train""": jsonl_path}						,  cache_dir=lowerCAmelCase						,  keep_in_memory=lowerCAmelCase ).read()
		_check_json_datasetdict(lowerCAmelCase						,  lowerCAmelCase )


@pytest.mark.parametrize(
    """features"""						,  [
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]						,  )
def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				Optional[Any]:

		"""simple docstring"""






		_UpperCamelCase						    =	tmp_path / """cache"""
		_UpperCamelCase						    =	{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
		_UpperCamelCase						    =	features.copy() if features else default_expected_features
		_UpperCamelCase						    =	(
		    Features({feature: Value(lowerCAmelCase ) for feature, dtype in features.items()} ) if features is not None else None
		)
		_UpperCamelCase						    =	JsonDatasetReader({"""train""": jsonl_path}						,  features=lowerCAmelCase						,  cache_dir=lowerCAmelCase ).read()
		_check_json_datasetdict(lowerCAmelCase						,  lowerCAmelCase )


@pytest.mark.parametrize("""split"""						,  [None, NamedSplit("""train""" ), """train""", """test"""] )
def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				int:

		"""simple docstring"""






		if split:
				_UpperCamelCase						    =	{split: jsonl_path}
		else:
				_UpperCamelCase						    =	"""train"""
				_UpperCamelCase						    =	{"""train""": jsonl_path, """test""": jsonl_path}
		_UpperCamelCase						    =	tmp_path / """cache"""
		_UpperCamelCase						    =	{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
		_UpperCamelCase						    =	JsonDatasetReader(lowerCAmelCase						,  cache_dir=lowerCAmelCase ).read()
		_check_json_datasetdict(lowerCAmelCase						,  lowerCAmelCase						,  splits=list(path.keys() ) )
		assert all(dataset[split].split == split for split in path.keys() )


def 							__A(lowerCAmelCase )		->				Optional[int]:

		"""simple docstring"""






		return json.load(lowerCAmelCase )


def 							__A(lowerCAmelCase )		->				Dict:

		"""simple docstring"""






		return [json.loads(lowerCAmelCase ) for line in buffer]







class        lowerCAmelCase__      :

						@pytest.mark.parametrize("""lines, load_json_function"""       ,							[(True, load_json_lines), (False, load_json)]				)
						def  A_				(			self       ,							a       ,							a       ,							a				)					->      List[str]:


								'''simple docstring'''

								with io.BytesIO() as buffer:
										JsonDatasetWriter(a       ,							a       ,							lines=a				).write()
										buffer.seek(0				)
										_UpperCamelCase						    =	load_json_function(a				)
								assert isinstance(a       ,							a				)
								assert isinstance(exported_content[0]       ,							a				)
								assert len(a				) == 10

						@pytest.mark.parametrize(
						    """orient, container, keys, len_at"""       ,							[
						        ("""records""", list, {"""tokens""", """labels""", """answers""", """id"""}, None),
						        ("""split""", dict, {"""columns""", """data"""}, """data"""),
						        ("""index""", dict, set("""0123456789"""				), None),
						        ("""columns""", dict, {"""tokens""", """labels""", """answers""", """id"""}, """tokens"""),
						        ("""values""", list, None, None),
						        ("""table""", dict, {"""schema""", """data"""}, """data"""),
						    ]       ,							)
						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Union[str, Any]:


								'''simple docstring'''

								with io.BytesIO() as buffer:
										JsonDatasetWriter(a       ,							a       ,							lines=a       ,							orient=a				).write()
										buffer.seek(0				)
										_UpperCamelCase						    =	load_json(a				)
								assert isinstance(a       ,							a				)
								if keys:
										if container is dict:
												assert exported_content.keys() == keys
										else:
												assert exported_content[0].keys() == keys
								else:
										assert not hasattr(a       ,							"""keys"""				) and not hasattr(exported_content[0]       ,							"""keys"""				)
								if len_at:
										assert len(exported_content[len_at]				) == 10
								else:
										assert len(a				) == 10

						@pytest.mark.parametrize("""lines, load_json_function"""       ,							[(True, load_json_lines), (False, load_json)]				)
						def  A_				(			self       ,							a       ,							a       ,							a				)					->      str:


								'''simple docstring'''

								with io.BytesIO() as buffer:
										JsonDatasetWriter(a       ,							a       ,							lines=a       ,							num_proc=2				).write()
										buffer.seek(0				)
										_UpperCamelCase						    =	load_json_function(a				)
								assert isinstance(a       ,							a				)
								assert isinstance(exported_content[0]       ,							a				)
								assert len(a				) == 10

						@pytest.mark.parametrize(
						    """orient, container, keys, len_at"""       ,							[
						        ("""records""", list, {"""tokens""", """labels""", """answers""", """id"""}, None),
						        ("""split""", dict, {"""columns""", """data"""}, """data"""),
						        ("""index""", dict, set("""0123456789"""				), None),
						        ("""columns""", dict, {"""tokens""", """labels""", """answers""", """id"""}, """tokens"""),
						        ("""values""", list, None, None),
						        ("""table""", dict, {"""schema""", """data"""}, """data"""),
						    ]       ,							)
						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Optional[int]:


								'''simple docstring'''

								with io.BytesIO() as buffer:
										JsonDatasetWriter(a       ,							a       ,							lines=a       ,							orient=a       ,							num_proc=2				).write()
										buffer.seek(0				)
										_UpperCamelCase						    =	load_json(a				)
								assert isinstance(a       ,							a				)
								if keys:
										if container is dict:
												assert exported_content.keys() == keys
										else:
												assert exported_content[0].keys() == keys
								else:
										assert not hasattr(a       ,							"""keys"""				) and not hasattr(exported_content[0]       ,							"""keys"""				)
								if len_at:
										assert len(exported_content[len_at]				) == 10
								else:
										assert len(a				) == 10

						def  A_				(			self       ,							a				)					->      str:


								'''simple docstring'''

								with pytest.raises(a				):
										with io.BytesIO() as buffer:
												JsonDatasetWriter(a       ,							a       ,							num_proc=0				)







						@pytest.mark.parametrize("""compression, extension"""       ,							[("""gzip""", """gz"""), ("""bz2""", """bz2"""), ("""xz""", """xz""")]				)
						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Any:


								'''simple docstring'''

								_UpperCamelCase						    =	tmp_path_factory.mktemp("""data"""				) / F'test.json.{extension}'
								_UpperCamelCase						    =	str(shared_datadir / F'test_file.json.{extension}'				)
								JsonDatasetWriter(a       ,							a       ,							compression=a				).write()

								with fsspec.open(a       ,							"""rb"""       ,							compression="""infer"""				) as f:
										_UpperCamelCase						    =	f.read()
								with fsspec.open(a       ,							"""rb"""       ,							compression="""infer"""				) as f:
										_UpperCamelCase						    =	f.read()
								assert exported_content == original_content





from math import isclose, sqrt


def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				tuple[float, float, float]:

		"""simple docstring"""






		_UpperCamelCase						    =	point_y / 4 / point_x
		_UpperCamelCase						    =	2 * normal_gradient / (1 + normal_gradient * normal_gradient)
		_UpperCamelCase						    =	(1 - normal_gradient * normal_gradient) / (
		    1 + normal_gradient * normal_gradient
		)
		_UpperCamelCase						    =	(sa - ca * incoming_gradient) / (ca + sa * incoming_gradient)

		# to find the next point, solve the simultaeneous equations:
		# y^2 + 4x^2 = 100
		# y - b = m * (x - a)
		# ==> A x^2 + B x + C = 0
		_UpperCamelCase						    =	outgoing_gradient**2 + 4
		_UpperCamelCase						    =	2 * outgoing_gradient * (point_y - outgoing_gradient * point_x)
		_UpperCamelCase						    =	(point_y - outgoing_gradient * point_x) ** 2 - 1_0_0

		_UpperCamelCase						    =	(
		    -linear_term - sqrt(linear_term**2 - 4 * quadratic_term * constant_term )
		) / (2 * quadratic_term)
		_UpperCamelCase						    =	(
		    -linear_term + sqrt(linear_term**2 - 4 * quadratic_term * constant_term )
		) / (2 * quadratic_term)

		# two solutions, one of which is our input point
		_UpperCamelCase						    =	x_minus if isclose(lowerCAmelCase						,  lowerCAmelCase ) else x_plus
		_UpperCamelCase						    =	point_y + outgoing_gradient * (next_x - point_x)

		return next_x, next_y, outgoing_gradient


def 							__A(lowerCAmelCase = 1.4						,  lowerCAmelCase = -9.6 )		->				int:

		"""simple docstring"""






		_UpperCamelCase						    =	0
		_UpperCamelCase						    =	first_x_coord
		_UpperCamelCase						    =	first_y_coord
		_UpperCamelCase						    =	(10.1 - point_y) / (0.0 - point_x)

		while not (-0.01 <= point_x <= 0.01 and point_y > 0):
				_UpperCamelCase    ,	_UpperCamelCase    ,	_UpperCamelCase						    =	next_point(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )
				num_reflections += 1

		return num_reflections


if __name__ == "__main__":
		print(F"""{solution() = }""")



