





"""simple docstring"""





# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
  import jax
  import jaxlib

__A											=		get_logger()

__A											=		None






class 					_lowerCAmelCase   (     TensorFormatter[Mapping, """jax.Array""", Mapping]       ):
       """simple docstring"""




       def __init__(    self    , __UpperCAmelCase=None    , __UpperCAmelCase=None    , **__UpperCAmelCase      ):



         '''simple docstring'''

         super().__init__(features=__UpperCAmelCase      )
         import jax
         from jaxlib.xla_client import Device

         if isinstance(__UpperCAmelCase    , __UpperCAmelCase      ):
           raise ValueError(
               F"Expected {device} to be a `str` not {type(__UpperCAmelCase      )}, as `jaxlib.xla_extension.Device` "
               'is not serializable neither with `pickle` nor with `dill`. Instead you can surround '
               'the device with `str()` to get its string identifier that will be internally mapped '
               'to the actual `jaxlib.xla_extension.Device`.'      )
         lowerCAmelCase__			:Optional[int]											=							device if isinstance(__UpperCAmelCase    , __UpperCAmelCase      ) else str(jax.devices()[0]      )
         # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
         # with `pickle` nor with `dill`, so we need to use a global variable instead
         global DEVICE_MAPPING
         if DEVICE_MAPPING is None:
           lowerCAmelCase__			:Union[str, Any]											=							self._map_devices_to_str()
         if self.device not in list(DEVICE_MAPPING.keys()      ):
           logger.warning(
               F"Device with string identifier {self.device} not listed among the available "
               F"devices: {list(DEVICE_MAPPING.keys()      )}, so falling back to the default "
               F"device: {str(jax.devices()[0]      )}."      )
           lowerCAmelCase__			:Optional[Any]											=							str(jax.devices()[0]      )
         lowerCAmelCase__			:Tuple											=							jnp_array_kwargs






       @staticmethod
       def      snake_case						(    ):



         '''simple docstring'''

         import jax

         return {str(__UpperCAmelCase      ): device for device in jax.devices()}






       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

         import jax
         import jax.numpy as jnp

         if isinstance(__UpperCAmelCase    , __UpperCAmelCase      ) and column:
           if all(
               isinstance(__UpperCAmelCase    , jax.Array      ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column      ):
             return jnp.stack(__UpperCAmelCase    , axis=0      )
         return column






       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

         import jax
         import jax.numpy as jnp

         if isinstance(__UpperCAmelCase    , (str, bytes, type(__UpperCAmelCase      ))      ):
           return value
         elif isinstance(__UpperCAmelCase    , (np.character, np.ndarray)      ) and np.issubdtype(value.dtype    , np.character      ):
           return value.tolist()

         lowerCAmelCase__			:int											=							{}

         if isinstance(__UpperCAmelCase    , (np.number, np.ndarray)      ) and np.issubdtype(value.dtype    , np.integer      ):
           # the default int precision depends on the jax config
           # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
           if jax.config.jax_enable_xaa:
             lowerCAmelCase__			:Dict											=							{'dtype': jnp.intaa}
           else:
             lowerCAmelCase__			:Any											=							{'dtype': jnp.intaa}
         elif isinstance(__UpperCAmelCase    , (np.number, np.ndarray)      ) and np.issubdtype(value.dtype    , np.floating      ):
           lowerCAmelCase__			:int											=							{'dtype': jnp.floataa}
         elif config.PIL_AVAILABLE and "PIL" in sys.modules:
           import PIL.Image

           if isinstance(__UpperCAmelCase    , PIL.Image.Image      ):
             lowerCAmelCase__			:Tuple											=							np.asarray(__UpperCAmelCase      )

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
         global DEVICE_MAPPING
         if DEVICE_MAPPING is None:
           lowerCAmelCase__			:Union[str, Any]											=							self._map_devices_to_str()

         with jax.default_device(DEVICE_MAPPING[self.device]      ):
           # calling jnp.array on a np.ndarray does copy the data
           # see https://github.com/google/jax/issues/4486
           return jnp.array(__UpperCAmelCase    , **{**default_dtype, **self.jnp_array_kwargs}      )






       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

         import jax

         # support for torch, tf, jax etc.
         if config.TORCH_AVAILABLE and "torch" in sys.modules:
           import torch

           if isinstance(__UpperCAmelCase    , torch.Tensor      ):
             return self._tensorize(data_struct.detach().cpu().numpy()[()]      )
         if hasattr(__UpperCAmelCase    , '__array__'      ) and not isinstance(__UpperCAmelCase    , jax.Array      ):
           lowerCAmelCase__			:List[str]											=							data_struct.__array__()
         # support for nested types like struct of list of struct
         if isinstance(__UpperCAmelCase    , np.ndarray      ):
           if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
             return self._consolidate([self.recursive_tensorize(__UpperCAmelCase      ) for substruct in data_struct]      )
         elif isinstance(__UpperCAmelCase    , (list, tuple)      ):
           return self._consolidate([self.recursive_tensorize(__UpperCAmelCase      ) for substruct in data_struct]      )
         return self._tensorize(__UpperCAmelCase      )






       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

         return map_nested(self._recursive_tensorize    , __UpperCAmelCase    , map_list=__UpperCAmelCase      )






       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

         lowerCAmelCase__			:Dict											=							self.numpy_arrow_extractor().extract_row(__UpperCAmelCase      )
         lowerCAmelCase__			:Union[str, Any]											=							self.python_features_decoder.decode_row(__UpperCAmelCase      )
         return self.recursive_tensorize(__UpperCAmelCase      )






       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

         lowerCAmelCase__			:Union[str, Any]											=							self.numpy_arrow_extractor().extract_column(__UpperCAmelCase      )
         lowerCAmelCase__			:Tuple											=							self.python_features_decoder.decode_column(__UpperCAmelCase    , pa_table.column_names[0]      )
         lowerCAmelCase__			:int											=							self.recursive_tensorize(__UpperCAmelCase      )
         lowerCAmelCase__			:List[str]											=							self._consolidate(__UpperCAmelCase      )
         return column






       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

         lowerCAmelCase__			:Union[str, Any]											=							self.numpy_arrow_extractor().extract_batch(__UpperCAmelCase      )
         lowerCAmelCase__			:Tuple											=							self.python_features_decoder.decode_batch(__UpperCAmelCase      )
         lowerCAmelCase__			:Optional[Any]											=							self.recursive_tensorize(__UpperCAmelCase      )
         for column_name in batch:
           lowerCAmelCase__			:Dict											=							self._consolidate(batch[column_name]      )
         return batch


import argparse
import hashlib  # hashlib is only used inside the Test class
import struct







class      A__							:






				def __init__(						self					:      Dict    ,     a					:      str   ):






							'''simple docstring'''







							lowerCAmelCase__      :      Optional[int]         =		data
							lowerCAmelCase__      :      List[Any]         =		[0X67_452_301, 0Xef_cda_b89, 0X98_bad_cfe, 0X10_325_476, 0Xc3_d2e_1f0]






				@staticmethod
				def 							_lowerCamelCase		(						a					:      int    ,     a					:      List[str]   ):






							'''simple docstring'''







							return ((n << b) | (n >> (32 - b))) & 0Xff_fff_fff






				def 							_lowerCamelCase		(						self					:      Optional[int]   ):






							'''simple docstring'''







							lowerCAmelCase__      :      Any         =		b'\x80' + b'\x00' * (63 - (len(self.data   ) + 8) % 64)
							lowerCAmelCase__      :      List[str]         =		self.data + padding + struct.pack('>Q'    ,     8 * len(self.data   )   )
							return padded_data






				def 							_lowerCamelCase		(						self					:      Tuple   ):






							'''simple docstring'''







							return [
							    self.padded_data[i : i + 64] for i in range(0    ,     len(self.padded_data   )    ,     64   )
							]






				def 							_lowerCamelCase		(						self					:      Tuple    ,     a					:      List[Any]   ):






							'''simple docstring'''







							lowerCAmelCase__      :      List[Any]         =		list(struct.unpack('>16L'    ,     a   )   ) + [0] * 64
							for i in range(16    ,     80   ):
										lowerCAmelCase__      :      int         =		self.rotate((w[i - 3] ^ w[i - 8] ^ w[i - 14] ^ w[i - 16])    ,     1   )
							return w






				def 							_lowerCamelCase		(						self					:      Dict   ):






							'''simple docstring'''







							lowerCAmelCase__      :      Optional[Any]         =		self.padding()
							lowerCAmelCase__      :      List[Any]         =		self.split_blocks()
							for block in self.blocks:
										lowerCAmelCase__      :      str         =		self.expand_block(a   )
										lowerCAmelCase__      ,				lowerCAmelCase__      ,				lowerCAmelCase__      ,				lowerCAmelCase__      ,				lowerCAmelCase__      :      str         =		self.h
										for i in range(0    ,     80   ):
													if 0 <= i < 20:
																lowerCAmelCase__      :      Tuple         =		(b & c) | ((~b) & d)
																lowerCAmelCase__      :      int         =		0X5a_827_999
													elif 20 <= i < 40:
																lowerCAmelCase__      :      List[str]         =		b ^ c ^ d
																lowerCAmelCase__      :      Any         =		0X6e_d9e_ba1
													elif 40 <= i < 60:
																lowerCAmelCase__      :      Tuple         =		(b & c) | (b & d) | (c & d)
																lowerCAmelCase__      :      Tuple         =		0X8f_1bb_cdc
													elif 60 <= i < 80:
																lowerCAmelCase__      :      List[Any]         =		b ^ c ^ d
																lowerCAmelCase__      :      int         =		0Xca_62c_1d6
													lowerCAmelCase__      ,				lowerCAmelCase__      ,				lowerCAmelCase__      ,				lowerCAmelCase__      ,				lowerCAmelCase__      :      Dict         =		(
													    self.rotate(a    ,     5   ) + f + e + k + expanded_block[i] & 0Xff_fff_fff,
													    a,
													    self.rotate(a    ,     30   ),
													    c,
													    d,
													)
										lowerCAmelCase__      :      Optional[Any]         =		(
										    self.h[0] + a & 0Xff_fff_fff,
										    self.h[1] + b & 0Xff_fff_fff,
										    self.h[2] + c & 0Xff_fff_fff,
										    self.h[3] + d & 0Xff_fff_fff,
										    self.h[4] + e & 0Xff_fff_fff,
										)
							return ("{:08x}" * 5).format(*self.h   )






def  lowerCAmelCase__			(  )				->		Tuple:
			lowerCAmelCase__      :      Optional[int]         =		b'Test String'
			assert SHAaHash(SCREAMING_SNAKE_CASE_     ).final_hash() == hashlib.shaa(SCREAMING_SNAKE_CASE_     ).hexdigest()  # noqa: S324
def  lowerCAmelCase__			(  )				->		Any:
			lowerCAmelCase__      :      str         =		argparse.ArgumentParser(description='Process some strings or files'     )
			parser.add_argument(
			    '--string'       ,							dest='input_string'       ,							default='Hello World!! Welcome to Cryptography'       ,							help='Hash the string'       ,							)
			parser.add_argument('--file'       ,							dest='input_file'       ,							help='Hash contents of a file'     )
			lowerCAmelCase__      :      Dict         =		parser.parse_args()
			lowerCAmelCase__      :      Tuple         =		args.input_string
			# In any case hash input should be a bytestring
			if args.input_file:
						with open(args.input_file       ,							'rb'     ) as f:
									lowerCAmelCase__      :      List[Any]         =		f.read()
			else:
						lowerCAmelCase__      :      Tuple         =		bytes(SCREAMING_SNAKE_CASE_       ,							'utf-8'     )
			print(SHAaHash(SCREAMING_SNAKE_CASE_     ).final_hash()     )


if __name__ == "__main__":
					main()
					import doctest

					doctest.testmod()