# Copyright (c) Facebook, Inc. and its affiliates.
#  
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import torch.nn as nn

import MinkowskiEngine as ME

from models.modules.common import ConvType, NormType
from models.modules.resnet_block import BasicBlock, Bottleneck


class SELayer(nn.Module):

  def __init__(self, channel, reduction=16, D=-1):
    # Global coords does not require coords_key
    super(SELayer, self).__init__()
    self.fc = nn.Sequential(
        ME.MinkowskiLinear(channel, channel // reduction), ME.MinkowskiReLU(inplace=True),
        ME.MinkowskiLinear(channel // reduction, channel), ME.MinkowskiSigmoid())
    self.pooling = ME.MinkowskiGlobalPooling(dimension=D)
    self.broadcast_mul = ME.MinkowskiBroadcastMultiplication(dimension=D)

  def forward(self, x):
    y = self.pooling(x)
    y = self.fc(y)
    return self.broadcast_mul(x, y)


class SEBasicBlock(BasicBlock):

  def __init__(self,
               inplanes,
               planes,
               stride=1,
               dilation=1,
               downsample=None,
               conv_type=ConvType.HYPERCUBE,
               reduction=16,
               D=-1):
    super(SEBasicBlock, self).__init__(
        inplanes,
        planes,
        stride=stride,
        dilation=dilation,
        downsample=downsample,
        conv_type=conv_type,
        D=D)
    self.se = SELayer(planes, reduction=reduction, D=D)

  def forward(self, x):
    residual = x

    out = self.conv1(x)
    out = self.norm1(out)
    out = self.relu(out)

    out = self.conv2(out)
    out = self.norm2(out)
    out = self.se(out)

    if self.downsample is not None:
      residual = self.downsample(x)

    out += residual
    out = self.relu(out)

    return out


class SEBasicBlockSN(SEBasicBlock):
  NORM_TYPE = NormType.SPARSE_SWITCH_NORM


class SEBasicBlockIN(SEBasicBlock):
  NORM_TYPE = NormType.SPARSE_INSTANCE_NORM


class SEBasicBlockLN(SEBasicBlock):
  NORM_TYPE = NormType.SPARSE_LAYER_NORM


class SEBottleneck(Bottleneck):

  def __init__(self,
               inplanes,
               planes,
               stride=1,
               dilation=1,
               downsample=None,
               conv_type=ConvType.HYPERCUBE,
               D=3,
               reduction=16):
    super(SEBottleneck, self).__init__(
        inplanes,
        planes,
        stride=stride,
        dilation=dilation,
        downsample=downsample,
        conv_type=conv_type,
        D=D)
    self.se = SELayer(planes * self.expansion, reduction=reduction, D=D)

  def forward(self, x):
    residual = x

    out = self.conv1(x)
    out = self.norm1(out)
    out = self.relu(out)

    out = self.conv2(out)
    out = self.norm2(out)
    out = self.relu(out)

    out = self.conv3(out)
    out = self.norm3(out)
    out = self.se(out)

    if self.downsample is not None:
      residual = self.downsample(x)

    out += residual
    out = self.relu(out)

    return out


class SEBottleneckSN(SEBottleneck):
  NORM_TYPE = NormType.SPARSE_SWITCH_NORM


class SEBottleneckIN(SEBottleneck):
  NORM_TYPE = NormType.SPARSE_INSTANCE_NORM


class SEBottleneckLN(SEBottleneck):
  NORM_TYPE = NormType.SPARSE_LAYER_NORM
