# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import torch
import torch.nn.functional as F
from torch import nn


class FPN(nn.Module):
    """
    Module that adds FPN on top of a list of feature maps.
    The feature maps are currently supposed to be in increasing depth
    order, and must be consecutive
    """

    def __init__(
        self,
        in_channels_list,
        out_channels,
        conv_block,
        top_blocks=None,
        drop_block=None,
        use_spp=False,
        use_pan=False,
        return_swint_feature_before_fusion=False,
    ):
        """
        Arguments:
            in_channels_list (list[int]): number of channels for each feature map that
                will be fed
            out_channels (int): number of channels of the FPN representation
            top_blocks (nn.Module or None): if provided, an extra operation will
                be performed on the output of the last (smallest resolution)
                FPN output, and the result will extend the result list
        """
        super(FPN, self).__init__()
        self.inner_blocks = []
        self.layer_blocks = []
        self.pan_blocks = [] if use_pan else None
        self.spp_block = SPPLayer() if use_spp else None
        self.return_swint_feature_before_fusion = return_swint_feature_before_fusion
        for idx, in_channels in enumerate(in_channels_list, 1):
            inner_block = "fpn_inner{}".format(idx)
            layer_block = "fpn_layer{}".format(idx)

            if in_channels == 0:
                continue
            if idx == len(in_channels_list) and use_spp:
                in_channels = in_channels * 4
            inner_block_module = conv_block(in_channels, out_channels, 1)
            layer_block_module = conv_block(out_channels, out_channels, 3, 1)
            self.add_module(inner_block, inner_block_module)
            self.add_module(layer_block, layer_block_module)
            self.inner_blocks.append(inner_block)
            self.layer_blocks.append(layer_block)

            if use_pan:
                pan_in_block = "pan_in_layer{}".format(idx)
                pan_in_block_module = conv_block(out_channels, out_channels, 3, 2)
                self.add_module(pan_in_block, pan_in_block_module)
                pan_out_block = "pan_out_layer{}".format(idx)
                pan_out_block_module = conv_block(out_channels, out_channels, 3, 1)
                self.add_module(pan_out_block, pan_out_block_module)
                self.pan_blocks.append([pan_in_block, pan_out_block])

        self.top_blocks = top_blocks
        self.drop_block = drop_block

    def forward(self, x):
        """
        Arguments:
            x (list[Tensor]): feature maps for each feature level.
        Returns:
            results (tuple[Tensor]): feature maps after FPN layers.
                They are ordered from highest resolution first.
        """
        if type(x) is tuple:
            # for the case of VL backbone
            x, x_text = x[0], x[1]
        # print([v.shape for v in x])
        swint_feature_c4 = None
        if self.return_swint_feature_before_fusion:
            # TODO: here we only return last single scale feature map before the backbone fusion, should be more flexible
            swint_feature_c4 = x[-2]

        if self.spp_block:
            last_inner = getattr(self, self.inner_blocks[-1])(self.spp_block(x[-1]))
        else:
            last_inner = getattr(self, self.inner_blocks[-1])(x[-1])
        results = []
        results.append(getattr(self, self.layer_blocks[-1])(last_inner))
        for feature, inner_block, layer_block in zip(
            x[:-1][::-1], self.inner_blocks[:-1][::-1], self.layer_blocks[:-1][::-1]
        ):
            if not inner_block:
                continue
            inner_lateral = getattr(self, inner_block)(feature)

            if inner_lateral.shape[-2:] != last_inner.shape[-2:]:
                # TODO: could also give size instead of
                inner_top_down = F.interpolate(last_inner, size=inner_lateral.shape[-2:], mode="nearest")
            else:
                inner_top_down = last_inner

            # TODO use size instead of scale to make it robust to different sizes
            # inner_top_down = F.upsample(last_inner, size=inner_lateral.shape[-2:],
            # mode='bilinear', align_corners=False)
            last_inner = inner_lateral + inner_top_down
            if self.drop_block and self.training:
                results.insert(0, getattr(self, layer_block)(self.drop_block(last_inner)))
            else:
                results.insert(0, getattr(self, layer_block)(last_inner))

        if self.pan_blocks:
            pan_results = []
            last_outer = results[0]
            pan_results.append(last_outer)
            for outer_top_down, pan_block in zip(results[1:], self.pan_blocks):

                if self.drop_block and self.training:
                    pan_lateral = getattr(self, pan_block[0])(self.drop_block(last_outer))
                else:
                    pan_lateral = getattr(self, pan_block[0])(last_outer)

                last_outer = getattr(self, pan_block[1])(pan_lateral + outer_top_down)
                pan_results.append(last_outer)
            results = pan_results

        if isinstance(self.top_blocks, LastLevelP6P7):
            last_results = self.top_blocks(x[-1], results[-1])
            results.extend(last_results)
        elif isinstance(self.top_blocks, LastLevelMaxPool):
            last_results = self.top_blocks(results[-1])
            results.extend(last_results)

        try:
            return tuple(results), x_text, swint_feature_c4
        except NameError as e:
            return tuple(results)


class LastLevelMaxPool(nn.Module):
    def forward(self, x):
        return [F.max_pool2d(x, 1, 2, 0)]


class LastLevelP6P7(nn.Module):
    """
    This module is used in RetinaNet to generate extra layers, P6 and P7.
    """

    def __init__(self, in_channels, out_channels):
        super(LastLevelP6P7, self).__init__()
        self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
        self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)
        for module in [self.p6, self.p7]:
            nn.init.kaiming_uniform_(module.weight, a=1)
            nn.init.constant_(module.bias, 0)
        self.use_P5 = in_channels == out_channels

    def forward(self, c5, p5):
        x = p5 if self.use_P5 else c5
        p6 = self.p6(x)
        p7 = self.p7(F.relu(p6))
        return [p6, p7]


class SPPLayer(nn.Module):
    def __init__(self):
        super(SPPLayer, self).__init__()

    def forward(self, x):
        x_1 = x
        x_2 = F.max_pool2d(x, 5, stride=1, padding=2)
        x_3 = F.max_pool2d(x, 9, stride=1, padding=4)
        x_4 = F.max_pool2d(x, 13, stride=1, padding=6)
        out = torch.cat((x_1, x_2, x_3, x_4), dim=1)
        return out