import torch
from typing import Callable


def encode_and_aggregate(input_tensor: torch.Tensor,
                         encoder: torch.nn.Module,
                         num_encoder_input_channels: int,
                         num_image_channels: int,
                         encode_channels_jointly: bool,
                         aggregation_layer: Callable) -> torch.Tensor:
    """
    Function that encodes a given input tensor either jointly using the encoder or separately for each channel
    in a sequential manner. Features obtained at the output encoder are then aggregated with the pooling function
    defined by `aggregation layer`.
    """
    if encode_channels_jointly:
        input_tensor = encoder(input_tensor)
        input_tensor = aggregation_layer(input_tensor)
    else:
        shape = input_tensor.shape
        channel_shape = (shape[0], num_encoder_input_channels, shape[2], shape[3], shape[4])
        encode_and_aggregate = []
        # When using multiple encoders, it is more memory efficient to aggregate the individual
        # encoder outputs and then stack those smaller results, rather than stack huge outputs and aggregate.
        for i in range(num_image_channels):
            start_index = i * num_encoder_input_channels
            end_index = start_index + num_encoder_input_channels
            encoder_output = encoder(input_tensor[:, start_index:end_index].view(channel_shape))
            aggregated = aggregation_layer(encoder_output)
            encode_and_aggregate.append(aggregated)
        input_tensor = torch.cat(encode_and_aggregate, dim=1)
    return input_tensor