
import inspect
import unittest

from transformers import DecisionTransformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch

			from transformers import DecisionTransformerModel
			from transformers.models.decision_transformer.modeling_decision_transformer import (
			    DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
			)




class     _SCREAMING_SNAKE_CASE     :



	def __init__(    self						,  lowercase						,  lowercase=13						,  lowercase=7						,  lowercase=6						,  lowercase=17						,  lowercase=23						,  lowercase=11						,  lowercase=True						,  )				->       int:
					lowerCamelCase_						=			parent
					lowerCamelCase_						=			batch_size
					lowerCamelCase_						=			seq_length
					lowerCamelCase_						=			act_dim
					lowerCamelCase_						=			state_dim
					lowerCamelCase_						=			hidden_size
					lowerCamelCase_						=			max_length
					lowerCamelCase_						=			is_training



	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					lowerCamelCase_						=			floats_tensor((self.batch_size, self.seq_length, self.state_dim) )
					lowerCamelCase_						=			floats_tensor((self.batch_size, self.seq_length, self.act_dim) )
					lowerCamelCase_						=			floats_tensor((self.batch_size, self.seq_length, 1) )
					lowerCamelCase_						=			floats_tensor((self.batch_size, self.seq_length, 1) )
					lowerCamelCase_						=			ids_tensor((self.batch_size, self.seq_length)						,  vocab_size=1000 )
					lowerCamelCase_						=			random_attention_mask((self.batch_size, self.seq_length) )

					lowerCamelCase_						=			self.get_config()

					return (
					    config,
					    states,
					    actions,
					    rewards,
					    returns_to_go,
					    timesteps,
					    attention_mask,
					)



	def     SCREAMING_SNAKE_CASE_(    self )				->       Any:
					return DecisionTransformerConfig(
					    batch_size=self.batch_size						,  seq_length=self.seq_length						,  act_dim=self.act_dim						,  state_dim=self.state_dim						,  hidden_size=self.hidden_size						,  max_length=self.max_length						,  )



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase						,  lowercase						,  lowercase						,  lowercase						,  lowercase						,  lowercase						,  lowercase						,  )				->       List[str]:
					lowerCamelCase_						=			DecisionTransformerModel(config=lowercase )
					model.to(lowercase )
					model.eval()
					lowerCamelCase_						=			model(lowercase						,  lowercase						,  lowercase						,  lowercase						,  lowercase						,  lowercase )

					self.parent.assertEqual(result.state_preds.shape						,  states.shape )
					self.parent.assertEqual(result.action_preds.shape						,  actions.shape )
					self.parent.assertEqual(result.return_preds.shape						,  returns_to_go.shape )
					self.parent.assertEqual(
					    result.last_hidden_state.shape						,  (self.batch_size, self.seq_length * 3, self.hidden_size) )  # seq length *3 as there are 3 modelities: states, returns and actions



	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[Any]:
					lowerCamelCase_						=			self.prepare_config_and_inputs()
					(
					    (
					    lowerCamelCase_
					)	,    (
					    lowerCamelCase_
					)	,    (
					    lowerCamelCase_
					)	,    (
					    lowerCamelCase_
					)	,    (
					    lowerCamelCase_
					)	,    (
					    lowerCamelCase_
					)	,    (
					    lowerCamelCase_
					)	,    
					)						=			config_and_inputs
					lowerCamelCase_						=			{
					    "states": states,
					    "actions": actions,
					    "rewards": rewards,
					    "returns_to_go": returns_to_go,
					    "timesteps": timesteps,
					    "attention_mask": attention_mask,
					}
					return config, inputs_dict




@require_torch
class     _SCREAMING_SNAKE_CASE     (      snake_case_			,       snake_case_			,       snake_case_			,       unittest.TestCase						):
	lowerCAmelCase__         = (DecisionTransformerModel,) if is_torch_available() else ()
	lowerCAmelCase__         = ()
	lowerCAmelCase__         = {'feature-extraction': DecisionTransformerModel} if is_torch_available() else {}

	# Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
	lowerCAmelCase__         = False

	# Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False



	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[int]:
					lowerCamelCase_						=			DecisionTransformerModelTester(self )
					lowerCamelCase_						=			ConfigTester(self						,  config_class=lowercase						,  hidden_size=37 )



	def     SCREAMING_SNAKE_CASE_(    self )				->       Any:
					self.config_tester.run_common_tests()



	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[int]:
					lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*lowercase )



	@slow
	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[Any]:
					for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowerCamelCase_						=			DecisionTransformerModel.from_pretrained(lowercase )
									self.assertIsNotNone(lowercase )



	def     SCREAMING_SNAKE_CASE_(    self )				->       Any:
					lowerCamelCase_	,    lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
									lowerCamelCase_						=			model_class(lowercase )
									lowerCamelCase_						=			inspect.signature(model.forward )
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									lowerCamelCase_						=			[*signature.parameters.keys()]

									lowerCamelCase_						=			[
									    "states",
									    "actions",
									    "rewards",
									    "returns_to_go",
									    "timesteps",
									    "attention_mask",
									]

									self.assertListEqual(arg_names[: len(lowercase )]						,  lowercase )




@require_torch
class     _SCREAMING_SNAKE_CASE     (      unittest.TestCase						):



	@slow
	def     SCREAMING_SNAKE_CASE_(    self )				->       List[Any]:
					lowerCamelCase_						=			2  # number of steps of autoregressive prediction we will perform
					lowerCamelCase_						=			10  # defined by the RL environment, may be normalized
					lowerCamelCase_						=			DecisionTransformerModel.from_pretrained("edbeeching/decision-transformer-gym-hopper-expert" )
					lowerCamelCase_						=			model.to(lowercase )
					lowerCamelCase_						=			model.config
					torch.manual_seed(0 )
					lowerCamelCase_						=			torch.randn(1						,  1						,  config.state_dim ).to(device=lowercase						,  dtype=torch.floataa )  # env.reset()

					lowerCamelCase_						=			torch.tensor(
					    [[0.2_4_2_7_9_3, -0.2_8_6_9_3_0_7_4, 0.8_7_4_2_6_1_3], [0.6_7_8_1_5_2_7_4, -0.0_8_1_0_1_0_8_5, -0.1_2_9_5_2_1_4_7]]						,  device=lowercase )

					lowerCamelCase_						=			torch.tensor(lowercase						,  device=lowercase						,  dtype=torch.floataa ).reshape(1						,  1						,  1 )
					lowerCamelCase_						=			state
					lowerCamelCase_						=			torch.zeros(1						,  0						,  config.act_dim						,  device=lowercase						,  dtype=torch.floataa )
					lowerCamelCase_						=			torch.zeros(1						,  0						,  device=lowercase						,  dtype=torch.floataa )
					lowerCamelCase_						=			torch.tensor(0						,  device=lowercase						,  dtype=torch.long ).reshape(1						,  1 )

					for step in range(lowercase ):
									lowerCamelCase_						=			torch.cat([actions, torch.zeros(1						,  1						,  config.act_dim						,  device=lowercase )]						,  dim=1 )
									lowerCamelCase_						=			torch.cat([rewards, torch.zeros(1						,  1						,  device=lowercase )]						,  dim=1 )

									lowerCamelCase_						=			torch.ones(1						,  states.shape[1] ).to(dtype=torch.long						,  device=states.device )

									with torch.no_grad():
													lowerCamelCase_	,    lowerCamelCase_	,    lowerCamelCase_						=			model(
													    states=lowercase						,  actions=lowercase						,  rewards=lowercase						,  returns_to_go=lowercase						,  timesteps=lowercase						,  attention_mask=lowercase						,  return_dict=lowercase						,  )

									self.assertEqual(action_pred.shape						,  actions.shape )
									self.assertTrue(torch.allclose(action_pred[0, -1]						,  expected_outputs[step]						,  atol=1e-4 ) )
									lowerCamelCase_	,    lowerCamelCase_	,    lowerCamelCase_	,    lowerCamelCase_						=			(  # env.step(action)
									    torch.randn(1						,  1						,  config.state_dim ).to(device=lowercase						,  dtype=torch.floataa ),
									    1.0,
									    False,
									    {},
									)

									lowerCamelCase_						=			action_pred[0, -1]
									lowerCamelCase_						=			torch.cat([states, state]						,  dim=1 )
									lowerCamelCase_						=			returns_to_go[0, -1] - reward
									lowerCamelCase_						=			torch.cat([returns_to_go, pred_return.reshape(1						,  1						,  1 )]						,  dim=1 )
									lowerCamelCase_						=			torch.cat(
									    [timesteps, torch.ones((1, 1)						,  device=lowercase						,  dtype=torch.long ) * (step + 1)]						,  dim=1 )




import math





def 			lowerCamelCase_ (				lowerCamelCase__	, lowerCamelCase__    ):
				if 0 not in (x, y):
								# We use the relation x^y = y*log10(x), where 10 is the base.
								return y * math.logaa(lowerCamelCase__    )
				else:
								if x == 0:  # 0 raised to any number is 0
												return 0
								elif y == 0:
												return 1  # any number raised to 0 is 1
				raise AssertionError("This should never happen"    )


if __name__ == "__main__":  # Main function
			# Read two numbers from input and typecast them to int using map function.
			# Here x is the base and y is the power.
			__A				='''Enter the base and the power separated by a comma: '''
			__A,						__A				=map(int, input(prompt).split(''','''))
			__A,						__A				=map(int, input(prompt).split(''','''))

			# We find the log of each number, using the function res(), which takes two
			# arguments.
			__A				=res(xa, ya)
			__A				=res(xa, ya)

			# We check for the largest number
			if resa > resa:
						print('''Largest number is''', xa, '''^''', ya)
			elif resa > resa:
						print('''Largest number is''', xa, '''^''', ya)
			else:
						print('''Both are equal''')


