from typing import Callable
from typing import List
from typing import Any


def get_images_batches(num_of_images: int, representative_data_gen: Callable) -> List[Any]:
    """
    Construct batches of image samples for inference.
    Args:
        num_of_images: Num of total images for evaluation.
        representative_data_gen: ataset used for getting batches for inference.
    Returns: A list of images batches (lists of images)
    """
    # First, select images to use for all measurements.
    samples_count = 0  # Number of images we used so far to compute the distance matrix.
    images_batches = []
    while samples_count < num_of_images:
        # Get a batch of images to infer in both models.
        inference_batch_input = representative_data_gen()
        batch_size = inference_batch_input[0].shape[0]

        # If we sampled more images than we should use in the distance matrix,
        # we take only a subset of these images and use only them for computing the distance matrix.
        if batch_size > num_of_images - samples_count:
            inference_batch_input = [x[:num_of_images - samples_count] for x in inference_batch_input]
            assert num_of_images - samples_count == inference_batch_input[0].shape[0]
            batch_size = num_of_images - samples_count

        images_batches.append(inference_batch_input)
        samples_count += batch_size
    return images_batches