#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

"""ResNe(X)t Head helper."""

import torch
import torch.nn as nn
# from detectron2.layers import ROIAlign


# class ResNetRoIHead(nn.Module):
#     """
#     ResNe(X)t RoI head.
#     """

#     def __init__(
#         self,
#         dim_in,
#         num_classes,
#         pool_size,
#         resolution,
#         scale_factor,
#         dropout_rate=0.0,
#         act_func="softmax",
#         aligned=True,
#     ):
#         """
#         The `__init__` method of any subclass should also contain these
#             arguments.
#         ResNetRoIHead takes p pathways as input where p in [1, infty].

#         Args:
#             dim_in (list): the list of channel dimensions of the p inputs to the
#                 ResNetHead.
#             num_classes (int): the channel dimensions of the p outputs to the
#                 ResNetHead.
#             pool_size (list): the list of kernel sizes of p spatial temporal
#                 poolings, temporal pool kernel size, spatial pool kernel size,
#                 spatial pool kernel size in order.
#             resolution (list): the list of spatial output size from the ROIAlign.
#             scale_factor (list): the list of ratio to the input boxes by this
#                 number.
#             dropout_rate (float): dropout rate. If equal to 0.0, perform no
#                 dropout.
#             act_func (string): activation function to use. 'softmax': applies
#                 softmax on the output. 'sigmoid': applies sigmoid on the output.
#             aligned (bool): if False, use the legacy implementation. If True,
#                 align the results more perfectly.
#         Note:
#             Given a continuous coordinate c, its two neighboring pixel indices
#             (in our pixel model) are computed by floor (c - 0.5) and ceil
#             (c - 0.5). For example, c=1.3 has pixel neighbors with discrete
#             indices [0] and [1] (which are sampled from the underlying signal at
#             continuous coordinates 0.5 and 1.5). But the original roi_align
#             (aligned=False) does not subtract the 0.5 when computing neighboring
#             pixel indices and therefore it uses pixels with a slightly incorrect
#             alignment (relative to our pixel model) when performing bilinear
#             interpolation.
#             With `aligned=True`, we first appropriately scale the ROI and then
#             shift it by -0.5 prior to calling roi_align. This produces the
#             correct neighbors; It makes negligible differences to the model's
#             performance if ROIAlign is used together with conv layers.
#         """
#         super(ResNetRoIHead, self).__init__()
#         assert (
#             len({len(pool_size), len(dim_in)}) == 1
#         ), "pathway dimensions are not consistent."
#         self.num_pathways = len(pool_size)
#         for pathway in range(self.num_pathways):
#             temporal_pool = nn.AvgPool3d(
#                 [pool_size[pathway][0], 1, 1], stride=1
#             )
#             self.add_module("s{}_tpool".format(pathway), temporal_pool)

#             roi_align = ROIAlign(
#                 resolution[pathway],
#                 spatial_scale=1.0 / scale_factor[pathway],
#                 sampling_ratio=0,
#                 aligned=aligned,
#             )
#             self.add_module("s{}_roi".format(pathway), roi_align)
#             spatial_pool = nn.MaxPool2d(resolution[pathway], stride=1)
#             self.add_module("s{}_spool".format(pathway), spatial_pool)

#         if dropout_rate > 0.0:
#             self.dropout = nn.Dropout(dropout_rate)

#         # Perform FC in a fully convolutional manner. The FC layer will be
#         # initialized with a different std comparing to convolutional layers.
#         self.projection = nn.Linear(sum(dim_in), num_classes, bias=True)

#         # Softmax for evaluation and testing.
#         if act_func == "softmax":
#             self.act = nn.Softmax(dim=1)
#         elif act_func == "sigmoid":
#             self.act = nn.Sigmoid()
#         else:
#             raise NotImplementedError(
#                 "{} is not supported as an activation"
#                 "function.".format(act_func)
#             )

#     def forward(self, inputs, bboxes):
#         assert (
#             len(inputs) == self.num_pathways
#         ), "Input tensor does not contain {} pathway".format(self.num_pathways)
#         pool_out = []
#         for pathway in range(self.num_pathways):
#             t_pool = getattr(self, "s{}_tpool".format(pathway))
#             out = t_pool(inputs[pathway])
#             assert out.shape[2] == 1
#             out = torch.squeeze(out, 2)

#             roi_align = getattr(self, "s{}_roi".format(pathway))
#             out = roi_align(out, bboxes)

#             s_pool = getattr(self, "s{}_spool".format(pathway))
#             pool_out.append(s_pool(out))

#         # B C H W.
#         x = torch.cat(pool_out, 1)

#         # Perform dropout.
#         if hasattr(self, "dropout"):
#             x = self.dropout(x)

#         x = x.view(x.shape[0], -1)
#         x = self.projection(x)
#         x = self.act(x)
#         return x


class ResNetBasicHead(nn.Module):
    """
    ResNe(X)t 3D head.
    This layer performs a fully-connected projection during training, when the
    input size is 1x1x1. It performs a convolutional projection during testing
    when the input size is larger than 1x1x1. If the inputs are from multiple
    different pathways, the inputs will be concatenated after pooling.
    """

    def __init__(
        self,
        dim_in,
        num_classes,
        pool_size,
        dropout_rate=0.0,
        act_func="softmax",
    ):
        """
        The `__init__` method of any subclass should also contain these
            arguments.
        ResNetBasicHead takes p pathways as input where p in [1, infty].

        Args:
            dim_in (list): the list of channel dimensions of the p inputs to the
                ResNetHead.
            num_classes (int): the channel dimensions of the p outputs to the
                ResNetHead.
            pool_size (list): the list of kernel sizes of p spatial temporal
                poolings, temporal pool kernel size, spatial pool kernel size,
                spatial pool kernel size in order.
            dropout_rate (float): dropout rate. If equal to 0.0, perform no
                dropout.
            act_func (string): activation function to use. 'softmax': applies
                softmax on the output. 'sigmoid': applies sigmoid on the output.
        """
        super(ResNetBasicHead, self).__init__()
        assert (
            len({len(pool_size), len(dim_in)}) == 1
        ), "pathway dimensions are not consistent."
        self.num_pathways = len(pool_size)

        for pathway in range(self.num_pathways):
            if pool_size[pathway] is None:
                avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
            else:
                avg_pool = nn.AvgPool3d(pool_size[pathway], stride=1)
            self.add_module("pathway{}_avgpool".format(pathway), avg_pool)

        if dropout_rate > 0.0:
            self.dropout = nn.Dropout(dropout_rate)
        # Perform FC in a fully convolutional manner. The FC layer will be
        # initialized with a different std comparing to convolutional layers.
        self.projection = nn.Linear(sum(dim_in), num_classes, bias=True)

        # Softmax for evaluation and testing.
        if act_func == "softmax":
            self.act = nn.Softmax(dim=4)
        elif act_func == "sigmoid":
            self.act = nn.Sigmoid()
        else:
            raise NotImplementedError(
                "{} is not supported as an activation"
                "function.".format(act_func)
            )

    def forward(self, inputs):
        assert (
            len(inputs) == self.num_pathways
        ), "Input tensor does not contain {} pathway".format(self.num_pathways)
        pool_out = []
        for pathway in range(self.num_pathways):
            m = getattr(self, "pathway{}_avgpool".format(pathway))
            pool_out.append(m(inputs[pathway]))
        x = torch.cat(pool_out, 1)
        # (N, C, T, H, W) -> (N, T, H, W, C).
        x = x.permute((0, 2, 3, 4, 1))
        # Perform dropout.
        if hasattr(self, "dropout"):
            x = self.dropout(x)
        x = self.projection(x)

        # Performs fully convlutional inference.
        if not self.training:
            x = self.act(x)
            x = x.mean([1, 2, 3])

        x = x.view(x.shape[0], -1)
        return x


class X3DHead(nn.Module):
    """
    X3D head.
    This layer performs a fully-connected projection during training, when the
    input size is 1x1x1. It performs a convolutional projection during testing
    when the input size is larger than 1x1x1. If the inputs are from multiple
    different pathways, the inputs will be concatenated after pooling.
    """

    def __init__(
        self,
        dim_in,
        dim_inner,
        dim_out,
        num_classes,
        pool_size,
        dropout_rate=0.0,
        act_func="softmax",
        inplace_relu=True,
        eps=1e-5,
        bn_mmt=0.1,
        norm_module=nn.BatchNorm3d,
        bn_lin5_on=False,
    ):
        """
        The `__init__` method of any subclass should also contain these
            arguments.
        X3DHead takes a 5-dim feature tensor (BxCxTxHxW) as input.

        Args:
            dim_in (float): the channel dimension C of the input.
            num_classes (int): the channel dimensions of the output.
            pool_size (float): a single entry list of kernel size for
                spatiotemporal pooling for the TxHxW dimensions.
            dropout_rate (float): dropout rate. If equal to 0.0, perform no
                dropout.
            act_func (string): activation function to use. 'softmax': applies
                softmax on the output. 'sigmoid': applies sigmoid on the output.
            inplace_relu (bool): if True, calculate the relu on the original
                input without allocating new memory.
            eps (float): epsilon for batch norm.
            bn_mmt (float): momentum for batch norm. Noted that BN momentum in
                PyTorch = 1 - BN momentum in Caffe2.
            norm_module (nn.Module): nn.Module for the normalization layer. The
                default is nn.BatchNorm3d.
            bn_lin5_on (bool): if True, perform normalization on the features
                before the classifier.
        """
        super(X3DHead, self).__init__()
        self.pool_size = pool_size
        self.dropout_rate = dropout_rate
        self.num_classes = num_classes
        self.act_func = act_func
        self.eps = eps
        self.bn_mmt = bn_mmt
        self.inplace_relu = inplace_relu
        self.bn_lin5_on = bn_lin5_on
        self._construct_head(dim_in, dim_inner, dim_out, norm_module)

    def _construct_head(self, dim_in, dim_inner, dim_out, norm_module):

        self.conv_5 = nn.Conv3d(
            dim_in,
            dim_inner,
            kernel_size=(1, 1, 1),
            stride=(1, 1, 1),
            padding=(0, 0, 0),
            bias=False,
        )
        self.conv_5_bn = norm_module(
            num_features=dim_inner, eps=self.eps, momentum=self.bn_mmt
        )
        self.conv_5_relu = nn.ReLU(self.inplace_relu)

        if self.pool_size is None:
            self.avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        else:
            self.avg_pool = nn.AvgPool3d(self.pool_size, stride=1)

        self.lin_5 = nn.Conv3d(
            dim_inner,
            dim_out,
            kernel_size=(1, 1, 1),
            stride=(1, 1, 1),
            padding=(0, 0, 0),
            bias=False,
        )
        if self.bn_lin5_on:
            self.lin_5_bn = norm_module(
                num_features=dim_out, eps=self.eps, momentum=self.bn_mmt
            )
        self.lin_5_relu = nn.ReLU(self.inplace_relu)

        if self.dropout_rate > 0.0:
            self.dropout = nn.Dropout(self.dropout_rate)
        # Perform FC in a fully convolutional manner. The FC layer will be
        # initialized with a different std comparing to convolutional layers.
        self.projection = nn.Linear(dim_out, self.num_classes, bias=True)

        # Softmax for evaluation and testing.
        if self.act_func == "softmax":
            self.act = nn.Softmax(dim=4)
        elif self.act_func == "sigmoid":
            self.act = nn.Sigmoid()
        else:
            raise NotImplementedError(
                "{} is not supported as an activation"
                "function.".format(self.act_func)
            )

    def forward(self, inputs):
        # In its current design the X3D head is only useable for a single
        # pathway input.
        assert len(inputs) == 1, "Input tensor does not contain 1 pathway"
        x = self.conv_5(inputs[0])
        x = self.conv_5_bn(x)
        x = self.conv_5_relu(x)
        x = self.avg_pool(x)

        x = self.lin_5(x)
        if self.bn_lin5_on:
            x = self.lin_5_bn(x)
        x = self.lin_5_relu(x)

        # (N, C, T, H, W) -> (N, T, H, W, C).
        x = x.permute((0, 2, 3, 4, 1))
        # Perform dropout.
        if hasattr(self, "dropout"):
            x = self.dropout(x)
        x = self.projection(x)

        # Performs fully convlutional inference.
        if not self.training:
            x = self.act(x)
            x = x.mean([1, 2, 3])

        x = x.view(x.shape[0], -1)
        return x


class TransformerBasicHead(nn.Module):
    """
    BasicHead. No pool.
    """

    def __init__(
        self,
        dim_in,
        num_classes,
        dropout_rate=0.0,
        act_func="softmax",
    ):
        """
        Perform linear projection and activation as head for tranformers.
        Args:
            dim_in (int): the channel dimension of the input to the head.
            num_classes (int): the channel dimensions of the output to the head.
            dropout_rate (float): dropout rate. If equal to 0.0, perform no
                dropout.
            act_func (string): activation function to use. 'softmax': applies
                softmax on the output. 'sigmoid': applies sigmoid on the output.
        """
        super(TransformerBasicHead, self).__init__()
        if dropout_rate > 0.0:
            self.dropout = nn.Dropout(dropout_rate)
        self.projection = nn.Linear(dim_in, num_classes, bias=True)

        # Softmax for evaluation and testing.
        if act_func == "softmax":
            self.act = nn.Softmax(dim=1)
        elif act_func == "sigmoid":
            self.act = nn.Sigmoid()
        else:
            raise NotImplementedError(
                "{} is not supported as an activation"
                "function.".format(act_func)
            )

    def forward(self, x):
        if hasattr(self, "dropout"):
            x = self.dropout(x)
        x = self.projection(x)

        if not self.training:
            x = self.act(x)
        return x
