



import random
import unittest

from torch.utils.data import BatchSampler, DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.data_loader import (
    BatchSamplerShard,
    DataLoaderDispatcher,
    DataLoaderShard,
    IterableDatasetShard,
    SkipBatchSampler,
    SkipDataLoader,
    skip_first_batches,
)


class 	lowerCamelCase__	(       _A):



 """simple docstring"""

 def __init__(    self					:					Dict			,						__lowerCAmelCase					:					List[str]=0.01			,						__lowerCAmelCase					:					Optional[Any]=10_00      )   ->	Tuple:
    _A      =       p_stop
    _A      =       max_length

 def __iter__(    self					:					Tuple      )   ->	int:
    _A      =       0
    _A      =       False
    while not stop and count < self.max_length:
       yield count
       count += 1
       _A      =       random.random() < self.p_stop


class 	lowerCamelCase__	(       unittest.TestCase):



 """simple docstring"""

 def     snake_case_      (    self					:					List[Any]			,						__lowerCAmelCase					:					Union[str, Any]			,						__lowerCAmelCase					:					List[str]			,						__lowerCAmelCase					:					List[Any]=False			,						__lowerCAmelCase					:					str=True      )   ->	Optional[int]:
    _A      =       [
        BatchSamplerShard(__lowerCAmelCase			,						2			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )
        for i in range(2      )
    ]
    _A      =       [list(__lowerCAmelCase      ) for batch_sampler_shard in batch_sampler_shards]
    if not split_batches:
       self.assertListEqual([len(__lowerCAmelCase      ) for shard in batch_sampler_shards]			,						[len(__lowerCAmelCase      ) for e in expected]      )
    self.assertListEqual(__lowerCAmelCase			,						__lowerCAmelCase      )

 def     snake_case_      (    self					:					Union[str, Any]      )   ->	Dict:
    # Check the shards when the dataset is a round multiple of total batch size.
    _A      =       BatchSampler(range(24      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase      )

    _A      =       BatchSampler(range(24      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    # Expected shouldn't change
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase      )

    # Check the shards when the dataset is a round multiple of batch size but not total batch size.
    _A      =       BatchSampler(range(21      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17], [0, 1, 2]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase      )

    _A      =       BatchSampler(range(21      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase      )

    # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
    # num_processes batch.
    _A      =       BatchSampler(range(22      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 0, 1]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase      )

    _A      =       BatchSampler(range(22      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase      )

    # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
    # num_processes batch.
    _A      =       BatchSampler(range(20      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 0]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17], [1, 2, 3]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase      )

    _A      =       BatchSampler(range(20      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase      )

    # Check the shards when the dataset is very small.
    _A      =       BatchSampler(range(2      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [[[0, 1, 0]], [[1, 0, 1]]]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase      )

    _A      =       BatchSampler(range(2      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [[], []]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase      )

 def     snake_case_      (    self					:					Optional[Any]      )   ->	Union[str, Any]:
    # Check the shards when the dataset is a round multiple of batch size.
    _A      =       BatchSampler(range(24      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
        [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(24      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    # Expected shouldn't change
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )

    # Check the shards when the dataset is not a round multiple of batch size.
    _A      =       BatchSampler(range(22      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
        [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [0, 1]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(22      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
        [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )

    # Check the shards when the dataset is not a round multiple of batch size or num_processes.
    _A      =       BatchSampler(range(21      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 0]],
        [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [1, 2]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(21      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
        [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )

    # Check the shards when the dataset is very small.
    _A      =       BatchSampler(range(2      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [[[0, 1]], [[0, 1]]]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(2      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [[], []]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )

 def     snake_case_      (    self					:					Optional[int]      )   ->	Union[str, Any]:
    # Check the shards when the dataset is a round multiple of total batch size.
    _A      =       BatchSampler(range(24      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(24      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    # Expected shouldn't change
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    # Check the shards when the dataset is a round multiple of batch size but not total batch size.
    _A      =       BatchSampler(range(21      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(21      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
    # num_processes batch.
    _A      =       BatchSampler(range(22      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(22      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
    # num_processes batch.
    _A      =       BatchSampler(range(20      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(20      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
        [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    # Check the shards when the dataset is very small.
    _A      =       BatchSampler(range(2      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [[[0, 1]], []]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(2      )			,						batch_size=3			,						drop_last=__lowerCAmelCase      )
    _A      =       [[], []]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

 def     snake_case_      (    self					:					str      )   ->	Dict:
    # Check the shards when the dataset is a round multiple of batch size.
    _A      =       BatchSampler(range(24      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
        [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(24      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    # Expected shouldn't change
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    # Check the shards when the dataset is not a round multiple of batch size.
    _A      =       BatchSampler(range(22      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
        [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(22      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
        [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    # Check the shards when the dataset is not a round multiple of batch size or num_processes.
    _A      =       BatchSampler(range(21      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
        [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(21      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [
        [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
        [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
    ]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    # Check the shards when the dataset is very small.
    _A      =       BatchSampler(range(2      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [[[0, 1]], []]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

    _A      =       BatchSampler(range(2      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       [[], []]
    self.check_batch_sampler_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						split_batches=__lowerCAmelCase			,						even_batches=__lowerCAmelCase      )

 def     snake_case_      (    self					:					List[Any]      )   ->	Any:
    _A      =       [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
    _A      =       [BatchSamplerShard(__lowerCAmelCase			,						2			,						__lowerCAmelCase			,						even_batches=__lowerCAmelCase      ) for i in range(2      )]

    self.assertEqual(len(batch_sampler_shards[0]      )			,						3      )
    self.assertEqual(len(batch_sampler_shards[1]      )			,						2      )

    self.assertListEqual(list(batch_sampler_shards[0]      )			,						[[0, 1, 2], [5, 6, 7, 8], [12, 13]]      )
    self.assertListEqual(list(batch_sampler_shards[1]      )			,						[[3, 4], [9, 10, 11]]      )

 def     snake_case_      (    self					:					Optional[int]			,						__lowerCAmelCase					:					Union[str, Any]			,						__lowerCAmelCase					:					Dict			,						__lowerCAmelCase					:					Tuple			,						__lowerCAmelCase					:					Union[str, Any]=False			,						__lowerCAmelCase					:					Any=2			,						__lowerCAmelCase					:					str=False      )   ->	Any:
    random.seed(__lowerCAmelCase      )
    _A      =       list(__lowerCAmelCase      )

    _A      =       [
        IterableDatasetShard(
            __lowerCAmelCase			,						batch_size=__lowerCAmelCase			,						drop_last=__lowerCAmelCase			,						num_processes=__lowerCAmelCase			,						process_index=__lowerCAmelCase			,						split_batches=__lowerCAmelCase			,						)
        for i in range(__lowerCAmelCase      )
    ]
    _A      =       []
    for iterable_dataset_shard in iterable_dataset_shards:
       # Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
       random.seed(__lowerCAmelCase      )
       iterable_dataset_lists.append(list(__lowerCAmelCase      )      )

    _A      =       batch_size // num_processes if split_batches else batch_size
    # All iterable dataset shard should have the same length, a round multiple of shard_batch_size
    _A      =       iterable_dataset_lists[0]
    for l in iterable_dataset_lists[1:]:
       self.assertEqual(len(__lowerCAmelCase      )			,						len(__lowerCAmelCase      )      )
       self.assertTrue(len(__lowerCAmelCase      ) % shard_batch_size == 0      )

    _A      =       []
    for idx in range(0			,						len(__lowerCAmelCase      )			,						__lowerCAmelCase      ):
       for l in iterable_dataset_lists:
          observed += l[idx : idx + shard_batch_size]

    if not drop_last:
       while len(__lowerCAmelCase      ) < len(__lowerCAmelCase      ):
          reference += reference
    self.assertListEqual(__lowerCAmelCase			,						reference[: len(__lowerCAmelCase      )]      )

 def     snake_case_      (    self					:					int      )   ->	str:
    _A      =       42
    _A      =       RandomIterableDataset()

    self.check_iterable_dataset_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						batch_size=4			,						drop_last=__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )
    self.check_iterable_dataset_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						batch_size=4			,						drop_last=__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )
    self.check_iterable_dataset_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						batch_size=4			,						drop_last=__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )
    self.check_iterable_dataset_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						batch_size=4			,						drop_last=__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )

    # Edge case with a very small dataset
    _A      =       RandomIterableDataset(max_length=2      )

    self.check_iterable_dataset_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						batch_size=4			,						drop_last=__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )
    self.check_iterable_dataset_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						batch_size=4			,						drop_last=__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )
    self.check_iterable_dataset_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						batch_size=4			,						drop_last=__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )
    self.check_iterable_dataset_shards(__lowerCAmelCase			,						__lowerCAmelCase			,						batch_size=4			,						drop_last=__lowerCAmelCase			,						split_batches=__lowerCAmelCase      )

 def     snake_case_      (    self					:					Dict      )   ->	str:
    _A      =       BatchSampler(range(16      )			,						batch_size=4			,						drop_last=__lowerCAmelCase      )
    _A      =       SkipBatchSampler(__lowerCAmelCase			,						2      )
    self.assertListEqual(list(__lowerCAmelCase      )			,						[[8, 9, 10, 11], [12, 13, 14, 15]]      )

 def     snake_case_      (    self					:					Dict      )   ->	List[str]:
    _A      =       SkipDataLoader(list(range(16      )      )			,						batch_size=4			,						skip_batches=2      )
    self.assertListEqual([t.tolist() for t in dataloader]			,						[[8, 9, 10, 11], [12, 13, 14, 15]]      )

 def     snake_case_      (    self					:					Any      )   ->	List[Any]:
    _A      =       DataLoader(list(range(16      )      )			,						batch_size=4      )
    _A      =       skip_first_batches(__lowerCAmelCase			,						num_batches=2      )
    self.assertListEqual([t.tolist() for t in new_dataloader]			,						[[8, 9, 10, 11], [12, 13, 14, 15]]      )

 def     snake_case_      (    self					:					Dict      )   ->	Any:
    _A      =       DataLoaderShard(list(range(16      )      )			,						batch_size=4      )
    for idx, _ in enumerate(__lowerCAmelCase      ):
       self.assertEqual(dataloader.end_of_dataloader			,						idx == 3      )

    # Test it also works on the second iteration
    for idx, _ in enumerate(__lowerCAmelCase      ):
       self.assertEqual(dataloader.end_of_dataloader			,						idx == 3      )

 def     snake_case_      (    self					:					List[Any]      )   ->	List[str]:
    Accelerator()
    _A      =       DataLoaderDispatcher(range(16      )			,						batch_size=4      )
    for idx, _ in enumerate(__lowerCAmelCase      ):
       self.assertEqual(dataloader.end_of_dataloader			,						idx == 3      )

    # Test it also works on the second iteration
    for idx, _ in enumerate(__lowerCAmelCase      ):
       self.assertEqual(dataloader.end_of_dataloader			,						idx == 3      )








from __future__ import annotations

from sys import maxsize
from typing import Generic, TypeVar

UpperCAmelCase_		 =					TypeVar("""T""")

def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:int				)  ->     int:
   return (position - 1) // 2

def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:int				)  ->     int:
   return (2 * position) + 1

def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:int				)  ->     int:
   return (2 * position) + 2


class 	lowerCamelCase__	(       Generic[T]):



 """simple docstring"""

 def __init__(    self					:					Optional[int]      )   ->	None:
    _A      =       []
    _A      =       {}
    _A      =       0

 def __len__(    self					:					str      )   ->	int:
    return self.elements

 def __repr__(    self					:					Optional[int]      )   ->	str:
    return str(self.heap      )

 def     snake_case_      (    self					:					str      )   ->	bool:
    # Check if the priority queue is empty
    return self.elements == 0

 def     snake_case_      (    self					:					Optional[int]			,						__lowerCAmelCase					:					T			,						__lowerCAmelCase					:					int      )   ->	None:
    # Add an element with given priority to the queue
    self.heap.append((elem, weight)      )
    _A      =       self.elements
    self.elements += 1
    self._bubble_up(__lowerCAmelCase      )

 def     snake_case_      (    self					:					Tuple      )   ->	T:
    # Remove and return the element with lowest weight (highest priority)
    if self.elements > 1:
       self._swap_nodes(0			,						self.elements - 1      )
    _A  ,    _A      =       self.heap.pop()
    del self.position_map[elem]
    self.elements -= 1
    if self.elements > 0:
       _A  ,    _A      =       self.heap[0]
       self._bubble_down(__lowerCAmelCase      )
    return elem

 def     snake_case_      (    self					:					int			,						__lowerCAmelCase					:					T			,						__lowerCAmelCase					:					int      )   ->	None:
    # Update the weight of the given key
    _A      =       self.position_map[elem]
    _A      =       (elem, weight)
    if position > 0:
       _A      =       get_parent_position(__lowerCAmelCase      )
       _A  ,    _A      =       self.heap[parent_position]
       if parent_weight > weight:
          self._bubble_up(__lowerCAmelCase      )
       else:
          self._bubble_down(__lowerCAmelCase      )
    else:
       self._bubble_down(__lowerCAmelCase      )

 def     snake_case_      (    self					:					Optional[Any]			,						__lowerCAmelCase					:					T      )   ->	None:
    # Place a node at the proper position (upward movement) [to be used internally
    # only]
    _A      =       self.position_map[elem]
    if curr_pos == 0:
       return None
    _A      =       get_parent_position(__lowerCAmelCase      )
    _A  ,    _A      =       self.heap[curr_pos]
    _A  ,    _A      =       self.heap[parent_position]
    if parent_weight > weight:
       self._swap_nodes(__lowerCAmelCase			,						__lowerCAmelCase      )
       return self._bubble_up(__lowerCAmelCase      )
    return None

 def     snake_case_      (    self					:					Dict			,						__lowerCAmelCase					:					T      )   ->	None:
    # Place a node at the proper position (downward movement) [to be used
    # internally only]
    _A      =       self.position_map[elem]
    _A  ,    _A      =       self.heap[curr_pos]
    _A      =       get_child_left_position(__lowerCAmelCase      )
    _A      =       get_child_right_position(__lowerCAmelCase      )
    if child_left_position < self.elements and child_right_position < self.elements:
       _A  ,    _A      =       self.heap[child_left_position]
       _A  ,    _A      =       self.heap[child_right_position]
       if child_right_weight < child_left_weight and child_right_weight < weight:
          self._swap_nodes(__lowerCAmelCase			,						__lowerCAmelCase      )
          return self._bubble_down(__lowerCAmelCase      )
    if child_left_position < self.elements:
       _A  ,    _A      =       self.heap[child_left_position]
       if child_left_weight < weight:
          self._swap_nodes(__lowerCAmelCase			,						__lowerCAmelCase      )
          return self._bubble_down(__lowerCAmelCase      )
    else:
       return None
    if child_right_position < self.elements:
       _A  ,    _A      =       self.heap[child_right_position]
       if child_right_weight < weight:
          self._swap_nodes(__lowerCAmelCase			,						__lowerCAmelCase      )
          return self._bubble_down(__lowerCAmelCase      )
    return None

 def     snake_case_      (    self					:					List[str]			,						__lowerCAmelCase					:					int			,						__lowerCAmelCase					:					int      )   ->	None:
    # Swap the nodes at the given positions
    _A      =       self.heap[nodea_pos][0]
    _A      =       self.heap[nodea_pos][0]
    _A  ,    _A      =       (
        self.heap[nodea_pos],
        self.heap[nodea_pos],
    )
    _A      =       nodea_pos
    _A      =       nodea_pos


class 	lowerCamelCase__	(       Generic[T]):



 """simple docstring"""

 def __init__(    self					:					str      )   ->	None:
    _A      =       {}
    _A      =       0

 def __repr__(    self					:					str      )   ->	str:
    return str(self.connections      )

 def __len__(    self					:					Dict      )   ->	int:
    return self.nodes

 def     snake_case_      (    self					:					Any			,						__lowerCAmelCase					:					T      )   ->	None:
    # Add a node in the graph if it is not in the graph
    if node not in self.connections:
       _A      =       {}
       self.nodes += 1

 def     snake_case_      (    self					:					str			,						__lowerCAmelCase					:					T			,						__lowerCAmelCase					:					T			,						__lowerCAmelCase					:					int      )   ->	None:
    # Add an edge between 2 nodes in the graph
    self.add_node(__lowerCAmelCase      )
    self.add_node(__lowerCAmelCase      )
    _A      =       weight
    _A      =       weight

def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:GraphUndirectedWeighted[T]	,   )  ->     tuple[dict[T, int], dict[T, T | None]]:
   _A      =       {node: maxsize for node in graph.connections}
   _A      =       {node: None for node in graph.connections}

   _A      =       MinPriorityQueue()
   for node, weight in dist.items():
      priority_queue.push(_snake_case	,   _snake_case				)

   if priority_queue.is_empty():
      return dist, parent

   # initialization
   _A      =       priority_queue.extract_min()
   _A      =       0
   for neighbour in graph.connections[node]:
      if dist[neighbour] > dist[node] + graph.connections[node][neighbour]:
         _A      =       dist[node] + graph.connections[node][neighbour]
         priority_queue.update_key(_snake_case	,   dist[neighbour]				)
         _A      =       node

    # running prim's algorithm
   while not priority_queue.is_empty():
      _A      =       priority_queue.extract_min()
      for neighbour in graph.connections[node]:
         if dist[neighbour] > dist[node] + graph.connections[node][neighbour]:
            _A      =       dist[node] + graph.connections[node][neighbour]
            priority_queue.update_key(_snake_case	,   dist[neighbour]				)
            _A      =       node
   return dist, parent



