import torch
from torch import nn


class MultiCropWrapper(nn.Module):
    """
    Perform forward pass separately on each resolution input.
    The inputs corresponding to a single resolution are clubbed and single
    forward is run on the same resolution inputs. Hence we do several
    forward passes = number of different resolutions used. We then
    concatenate all the output features and run the head forward on these
    concatenated features.
    """

    def __init__(
        self,
        backbone: torch.nn.Module,
        head: torch.nn.Module,
        apply_l2_norm: float = True,
    ):
        super(MultiCropWrapper, self).__init__()
        self.apply_l2_norm = apply_l2_norm
        # disable layers dedicated to ImageNet labels classification
        backbone.fc, backbone.head = nn.Identity(), nn.Identity()
        self.backbone = backbone
        if head is None:
            self.head = nn.Identity()
        else:
            self.head = head

    def forward(self, x, mask=None, return_backbone_feat=False, **kwargs):
        # convert to list
        if not isinstance(x, list):
            x = [x]
            mask = [mask] if mask is not None else None
        idx_crops = torch.cumsum(
            torch.unique_consecutive(
                torch.tensor([inp.shape[-1] for inp in x]),
                return_counts=True,
            )[1],
            0,
        )

        # start concatenating the features
        start_idx = 0
        output = torch.empty(0).to(x[0].device)
        for end_idx in idx_crops:
            inp_x = torch.cat(x[start_idx:end_idx])

            if mask is not None:
                inp_m = torch.cat(mask[start_idx:end_idx])
                kwargs.update(dict(mask=inp_m))

            _out = self.backbone(inp_x, **kwargs)
            if start_idx == 0:
                output = _out
            else:
                output = torch.cat((output, _out))
            start_idx = end_idx

        # SelfClean: apply L2 norm if requested
        if self.apply_l2_norm:
            output = torch.nn.functional.normalize(output, dim=-1, p=2)

        # run the head forward on the concatenated features
        output_ = self.head(output)
        if return_backbone_feat:
            return output, output_

        return output_
