



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
						import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


__A       					=		get_logger(__name__)



def  __a  (					lowerCAmelCase_     :      Any    ,lowerCAmelCase_     :      str    ,lowerCAmelCase_     :      List[str]    ,lowerCAmelCase_     :      str    ,lowerCAmelCase_     :      int    ,lowerCAmelCase_     :      List[str] = None    ,)   ->				Tuple:






						'''simple docstring'''



						for i, state in enumerate(lowerCamelCase__       ):
												UpperCAmelCase_=     F"""{MODEL_NAME}.bin""" if i == 0 else F"""{MODEL_NAME}_{i}.bin"""
												UpperCAmelCase_=     os.path.join(lowerCamelCase__    ,lowerCamelCase__       )
												save(lowerCamelCase__    ,lowerCamelCase__       )
												logger.info(F"""Model weights saved in {output_model_file}"""       )
						# Optimizer states
						for i, opt in enumerate(lowerCamelCase__       ):
												UpperCAmelCase_=     opt.state_dict()
												UpperCAmelCase_=     F"""{OPTIMIZER_NAME}.bin""" if i == 0 else F"""{OPTIMIZER_NAME}_{i}.bin"""
												UpperCAmelCase_=     os.path.join(lowerCamelCase__    ,lowerCamelCase__       )
												save(lowerCamelCase__    ,lowerCamelCase__       )
												logger.info(F"""Optimizer state saved in {output_optimizer_file}"""       )
						# Scheduler states
						for i, scheduler in enumerate(lowerCamelCase__       ):
												UpperCAmelCase_=     scheduler.state_dict()
												UpperCAmelCase_=     F"""{SCHEDULER_NAME}.bin""" if i == 0 else F"""{SCHEDULER_NAME}_{i}.bin"""
												UpperCAmelCase_=     os.path.join(lowerCamelCase__    ,lowerCamelCase__       )
												save(lowerCamelCase__    ,lowerCamelCase__       )
												logger.info(F"""Scheduler state saved in {output_scheduler_file}"""       )
						# GradScaler state
						if scaler is not None:
												UpperCAmelCase_=     scaler.state_dict()
												UpperCAmelCase_=     os.path.join(lowerCamelCase__    ,lowerCamelCase__       )
												torch.save(lowerCamelCase__    ,lowerCamelCase__       )
												logger.info(F"""Gradient scaler state saved in {output_scaler_file}"""       )
						# Random number generator states
						UpperCAmelCase_=     {}
						UpperCAmelCase_=     F"""{RNG_STATE_NAME}_{process_index}.pkl"""
						UpperCAmelCase_=     random.getstate()
						UpperCAmelCase_=     np.random.get_state()
						UpperCAmelCase_=     torch.get_rng_state()
						if is_xpu_available():
												UpperCAmelCase_=     torch.xpu.get_rng_state_all()
						else:
												UpperCAmelCase_=     torch.cuda.get_rng_state_all()
						if is_tpu_available():
												UpperCAmelCase_=     xm.get_rng_state()
						UpperCAmelCase_=     os.path.join(lowerCamelCase__    ,lowerCamelCase__       )
						torch.save(lowerCamelCase__    ,lowerCamelCase__       )
						logger.info(F"""Random states saved in {output_states_file}"""       )
						return output_dir



def  __a  (					lowerCAmelCase_     :      Optional[int]    ,lowerCAmelCase_     :      Tuple    ,lowerCAmelCase_     :      List[str]    ,lowerCAmelCase_     :      Union[str, Any]    ,lowerCAmelCase_     :      int    ,lowerCAmelCase_     :      Dict=None    ,lowerCAmelCase_     :      Union[str, Any]=None    ,**lowerCAmelCase_     :      Tuple    ,)   ->				List[str]:






						'''simple docstring'''



						if map_location not in [None, "cpu", "on_device"]:
												raise TypeError(
												    """Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`"""       )
						if map_location is None:
												UpperCAmelCase_=     """cpu"""
						elif map_location == "on_device":
												UpperCAmelCase_=     PartialState().device
						# Model states
						for i, model in enumerate(lowerCamelCase__       ):
												UpperCAmelCase_=     F"""{MODEL_NAME}.bin""" if i == 0 else F"""{MODEL_NAME}_{i}.bin"""
												UpperCAmelCase_=     os.path.join(lowerCamelCase__    ,lowerCamelCase__       )
												models[i].load_state_dict(torch.load(lowerCamelCase__    ,map_location=lowerCamelCase__       )    ,**lowerCamelCase__       )
						logger.info("""All model weights loaded successfully"""       )

						# Optimizer states
						for i, opt in enumerate(lowerCamelCase__       ):
												UpperCAmelCase_=     F"""{OPTIMIZER_NAME}.bin""" if i == 0 else F"""{OPTIMIZER_NAME}_{i}.bin"""
												UpperCAmelCase_=     os.path.join(lowerCamelCase__    ,lowerCamelCase__       )
												UpperCAmelCase_=     torch.load(lowerCamelCase__    ,map_location=lowerCamelCase__       )
												optimizers[i].load_state_dict(lowerCamelCase__       )
						logger.info("""All optimizer states loaded successfully"""       )

						# Scheduler states
						for i, scheduler in enumerate(lowerCamelCase__       ):
												UpperCAmelCase_=     F"""{SCHEDULER_NAME}.bin""" if i == 0 else F"""{SCHEDULER_NAME}_{i}.bin"""
												UpperCAmelCase_=     os.path.join(lowerCamelCase__    ,lowerCamelCase__       )
												scheduler.load_state_dict(torch.load(lowerCamelCase__       )       )
						logger.info("""All scheduler states loaded successfully"""       )

						# GradScaler state
						if scaler is not None:
												UpperCAmelCase_=     os.path.join(lowerCamelCase__    ,lowerCamelCase__       )
												scaler.load_state_dict(torch.load(lowerCamelCase__       )       )
												logger.info("""GradScaler state loaded successfully"""       )

						# Random states
						try:
												UpperCAmelCase_=     torch.load(os.path.join(lowerCamelCase__    ,F"""{RNG_STATE_NAME}_{process_index}.pkl"""       )       )
												random.setstate(states["""random_state"""]       )
												np.random.set_state(states["""numpy_random_seed"""]       )
												torch.set_rng_state(states["""torch_manual_seed"""]       )
												if is_xpu_available():
																		torch.xpu.set_rng_state_all(states["""torch_xpu_manual_seed"""]       )
												else:
																		torch.cuda.set_rng_state_all(states["""torch_cuda_manual_seed"""]       )
												if is_tpu_available():
																		xm.set_rng_state(states["""xm_seed"""]       )
												logger.info("""All random states loaded successfully"""       )
						except Exception:
												logger.info("""Could not load random states"""       )



def  __a  (					lowerCAmelCase_     :      int    ,lowerCAmelCase_     :      int    ,lowerCAmelCase_     :      Dict = 0       )   ->				Optional[int]:






						'''simple docstring'''



						UpperCAmelCase_=     Path(lowerCamelCase__       ) / F"""custom_checkpoint_{index}.pkl"""
						logger.info(F"""Saving the state of {get_pretty_name(lowerCamelCase__       )} to {save_location}"""       )
						torch.save(obj.state_dict()    ,lowerCamelCase__       )



def  __a  (					lowerCAmelCase_     :      List[str]    ,lowerCAmelCase_     :      Union[str, Any]    ,lowerCAmelCase_     :      Union[str, Any] = 0       )   ->				List[str]:






						'''simple docstring'''



						UpperCAmelCase_=     F"""{path}/custom_checkpoint_{index}.pkl"""
						logger.info(F"""Loading the state of {get_pretty_name(lowerCamelCase__       )} from {load_location}"""       )
						obj.load_state_dict(torch.load(lowerCamelCase__    ,map_location="""cpu"""       )       )









from math import radians as angle_to_radians
from math import sin

# Acceleration Constant on Earth (unit m/s^2)
__A       					=		9.8_0_6_6_5



def  __a  (					lowerCAmelCase_     :      float    ,lowerCAmelCase_     :      float       )   ->				None:






						'''simple docstring'''



						if not isinstance(lowerCAmelCase_    ,(int, float)       ):
												raise TypeError("""Invalid velocity. Should be a positive number."""       )

						if not isinstance(lowerCAmelCase_    ,(int, float)       ):
												raise TypeError("""Invalid angle. Range is 1-90 degrees."""       )

						# Ensure valid angle
						if angle > 90 or angle < 1:
												raise ValueError("""Invalid angle. Range is 1-90 degrees."""       )

						# Ensure valid velocity
						if init_velocity < 0:
												raise ValueError("""Invalid velocity. Should be a positive number."""       )



def  __a  (					lowerCAmelCase_     :      float    ,lowerCAmelCase_     :      float       )   ->				float:






						'''simple docstring'''



						check_args(lowerCAmelCase_    ,lowerCAmelCase_       )
						UpperCAmelCase_=     angle_to_radians(2 * angle       )
						return round(init_velocity**2 * sin(lowerCAmelCase_       ) / g    ,2       )



def  __a  (					lowerCAmelCase_     :      float    ,lowerCAmelCase_     :      float       )   ->				float:






						'''simple docstring'''



						check_args(lowerCAmelCase_    ,lowerCAmelCase_       )
						UpperCAmelCase_=     angle_to_radians(lowerCAmelCase_       )
						return round(init_velocity**2 * sin(lowerCAmelCase_       ) ** 2 / (2 * g)    ,2       )



def  __a  (					lowerCAmelCase_     :      float    ,lowerCAmelCase_     :      float       )   ->				float:






						'''simple docstring'''



						check_args(lowerCAmelCase_    ,lowerCAmelCase_       )
						UpperCAmelCase_=     angle_to_radians(lowerCAmelCase_       )
						return round(2 * init_velocity * sin(lowerCAmelCase_       ) / g    ,2       )



def  __a  (					)   ->				None:






						'''simple docstring'''



						UpperCAmelCase_,     UpperCAmelCase_=     25, 20
						assert horizontal_distance(lowerCAmelCase_    ,lowerCAmelCase_       ) == 40.97
						assert max_height(lowerCAmelCase_    ,lowerCAmelCase_       ) == 3.73
						assert total_time(lowerCAmelCase_    ,lowerCAmelCase_       ) == 1.74


if __name__ == "__main__":
	from doctest import testmod

	testmod()

	# Get input from user
	__A       					=		float(input('''Initial Velocity: ''').strip())

	# Get input from user
	__A       					=		float(input('''angle: ''').strip())

	# Print results
	print()
	print('''Results: ''')
	print(f'Horizontal Distance: {horizontal_distance(init_vel, angle)!s} [m]')
	print(f'Maximum Height: {max_height(init_vel, angle)!s} [m]')
	print(f'Total Time: {total_time(init_vel, angle)!s} [s]')




