


import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
__snake_case            =   object()

# For specifying empty leaf dict `{}`
__snake_case            =   object()






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		Union[str, Any]       ,  SCREAMING_SNAKE_CASE_   :		Dict	):





			"""simple docstring"""




			UpperCamelCase         =					tuple((re.compile(x + """$"""	) for x in qs)	)
			for i in range(len(SCREAMING_SNAKE_CASE_	) - len(SCREAMING_SNAKE_CASE_	) + 1	):
						UpperCamelCase         =					[x.match(SCREAMING_SNAKE_CASE_	) for x, y in zip(SCREAMING_SNAKE_CASE_       ,  ks[i:]	)]
						if matches and all(SCREAMING_SNAKE_CASE_	):
									return True
			return False






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		Optional[Any]	):





			"""simple docstring"""




			def replace(SCREAMING_SNAKE_CASE_   :		str       ,  SCREAMING_SNAKE_CASE_   :		int	):
						for rule, replacement in rules:
									if _match(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	):
												return replacement
						return val

			return replace






def 							_lowercase    (		):





			"""simple docstring"""




			return [
			    # embeddings
			    (("transformer", "wpe", "embedding"), P("""mp"""       ,  SCREAMING_SNAKE_CASE_	)),
			    (("transformer", "wte", "embedding"), P("""mp"""       ,  SCREAMING_SNAKE_CASE_	)),
			    # atention
			    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(SCREAMING_SNAKE_CASE_       ,  """mp"""	)),
			    (("attention", "out_proj", "kernel"), P("""mp"""       ,  SCREAMING_SNAKE_CASE_	)),
			    (("attention", "out_proj", "bias"), None),
			    # mlp
			    (("mlp", "c_fc", "kernel"), P(SCREAMING_SNAKE_CASE_       ,  """mp"""	)),
			    (("mlp", "c_fc", "bias"), P("""mp"""	)),
			    (("mlp", "c_proj", "kernel"), P("""mp"""       ,  SCREAMING_SNAKE_CASE_	)),
			    (("mlp", "c_proj", "bias"), None),
			    # layer norms
			    ((r"ln_\d+", "bias"), None),
			    ((r"\d+", r"ln_\d+", "scale"), None),
			    (("ln_f", "bias"), None),
			    (("ln_f", "scale"), None),
			]






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		Tuple	):





			"""simple docstring"""




			UpperCamelCase         =					_get_partition_rules()
			UpperCamelCase         =					_replacement_rules(SCREAMING_SNAKE_CASE_	)
			UpperCamelCase         =					{k: _unmatched for k in flatten_dict(SCREAMING_SNAKE_CASE_	)}
			UpperCamelCase         =					{k: replace(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	) for k, v in initd.items()}
			assert _unmatched not in result.values(), "Incomplete partition spec."
			return freeze(unflatten_dict(SCREAMING_SNAKE_CASE_	)	)








import os
import random
import sys

from . import cryptomath_module as cryptomath
from . import rabin_miller

__snake_case            =   3






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




			print("""Generating primitive root of p"""	)
			while True:
						UpperCamelCase         =					random.randrange(3       ,  SCREAMING_SNAKE_CASE_	)
						if pow(SCREAMING_SNAKE_CASE_       ,  2       ,  SCREAMING_SNAKE_CASE_	) == 1:
									continue
						if pow(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	) == 1:
									continue
						return g






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




			print("""Generating prime p..."""	)
			UpperCamelCase         =					rabin_miller.generate_large_prime(SCREAMING_SNAKE_CASE_	)  # select large prime number.
			UpperCamelCase         =					primitive_root(SCREAMING_SNAKE_CASE_	)  # one primitive root on modulo p.
			UpperCamelCase         =					random.randrange(3       ,  SCREAMING_SNAKE_CASE_	)  # private_key -> have to be greater than 2 for safety.
			UpperCamelCase         =					cryptomath.find_mod_inverse(pow(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)       ,  SCREAMING_SNAKE_CASE_	)

			UpperCamelCase         =					(key_size, e_a, e_a, p)
			UpperCamelCase         =					(key_size, d)

			return public_key, private_key






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		str       ,  SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




			if os.path.exists(f'{name}_pubkey.txt'	) or os.path.exists(f'{name}_privkey.txt'	):
						print("""\nWARNING:"""	)
						print(
						    f'"{name}_pubkey.txt" or "{name}_privkey.txt" already exists. \n'
						    """Use a different name or delete these files and re-run this program."""	)
						sys.exit()

			UpperCamelCase  ,  UpperCamelCase         =					generate_key(SCREAMING_SNAKE_CASE_	)
			print(f'\nWriting public key to file {name}_pubkey.txt...'	)
			with open(f'{name}_pubkey.txt'       ,  """w"""	) as fo:
						fo.write(f'{public_key[0]},{public_key[1]},{public_key[2]},{public_key[3]}'	)

			print(f'Writing private key to file {name}_privkey.txt...'	)
			with open(f'{name}_privkey.txt'       ,  """w"""	) as fo:
						fo.write(f'{private_key[0]},{private_key[1]}'	)






def 							_lowercase    (		):





			"""simple docstring"""




			print("""Making key files..."""	)
			make_key_files("""elgamal"""       ,  2_048	)
			print("""Key files generation successful"""	)


if __name__ == "__main__":
			main()




