"""MobileNet and MobileNetV2."""

"""
Code adopted from https://github.com/LikeLy-Journey/SegmenTron/blob/master/segmentron/models/backbones/mobilenet.py
"""
import torch
import torch.nn as nn
import torch.nn.functional as F

# ============  Basic Blocks  ============


class _ConvBNReLU(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        padding=0,
        dilation=1,
        groups=1,
        relu6=False,
        norm_layer=nn.BatchNorm2d,
    ):
        super(_ConvBNReLU, self).__init__()
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
            dilation,
            groups,
            bias=False,
        )
        self.bn = norm_layer(out_channels)
        self.relu = nn.ReLU6(True) if relu6 else nn.ReLU(True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x


class _DepthwiseConv(nn.Module):
    """conv_dw in MobileNet"""

    def __init__(
        self, in_channels, out_channels, stride, norm_layer=nn.BatchNorm2d, **kwargs
    ):
        super(_DepthwiseConv, self).__init__()
        self.conv = nn.Sequential(
            _ConvBNReLU(
                in_channels,
                in_channels,
                3,
                stride,
                1,
                groups=in_channels,
                norm_layer=norm_layer,
            ),
            _ConvBNReLU(in_channels, out_channels, 1, norm_layer=norm_layer),
        )

    def forward(self, x):
        return self.conv(x)


class InvertedResidual(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        stride,
        expand_ratio,
        dilation=1,
        norm_layer=nn.BatchNorm2d,
    ):
        super(InvertedResidual, self).__init__()
        assert stride in [1, 2]
        self.use_res_connect = stride == 1 and in_channels == out_channels

        layers = list()
        inter_channels = int(round(in_channels * expand_ratio))
        if expand_ratio != 1:
            # pw
            layers.append(
                _ConvBNReLU(
                    in_channels, inter_channels, 1, relu6=True, norm_layer=norm_layer
                )
            )
        layers.extend(
            [
                # dw
                _ConvBNReLU(
                    inter_channels,
                    inter_channels,
                    3,
                    stride,
                    dilation,
                    dilation,
                    groups=inter_channels,
                    relu6=True,
                    norm_layer=norm_layer,
                ),
                # pw-linear
                nn.Conv2d(inter_channels, out_channels, 1, bias=False),
                norm_layer(out_channels),
            ]
        )
        self.conv = nn.Sequential(*layers)

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


# ============  Backbone  ============


class MobileNetV2(nn.Module):
    def __init__(self, num_classes=1000, norm_layer=nn.BatchNorm2d):
        super(MobileNetV2, self).__init__()
        output_stride = 8
        self.multiplier = 1
        if output_stride == 32:
            dilations = [1, 1]
        elif output_stride == 16:
            dilations = [1, 2]
        elif output_stride == 8:
            dilations = [2, 4]
        else:
            raise NotImplementedError
        inverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],
            [6, 24, 2, 2],
            [6, 32, 3, 2],
            [6, 64, 4, 2],
            [6, 96, 3, 1],
            [6, 160, 3, 2],
            [6, 320, 1, 1],
        ]
        # building first layer
        input_channels = int(32 * self.multiplier) if self.multiplier > 1.0 else 32
        # last_channels = int(1280 * multiplier) if multiplier > 1.0 else 1280
        self.conv1 = _ConvBNReLU(
            3, input_channels, 3, 2, 1, relu6=True, norm_layer=norm_layer
        )

        # building inverted residual blocks
        self.planes = input_channels
        self.block1 = self._make_layer(
            InvertedResidual,
            self.planes,
            inverted_residual_setting[0:1],
            norm_layer=norm_layer,
        )
        self.block2 = self._make_layer(
            InvertedResidual,
            self.planes,
            inverted_residual_setting[1:2],
            norm_layer=norm_layer,
        )
        self.block3 = self._make_layer(
            InvertedResidual,
            self.planes,
            inverted_residual_setting[2:3],
            norm_layer=norm_layer,
        )
        self.block4 = self._make_layer(
            InvertedResidual,
            self.planes,
            inverted_residual_setting[3:5],
            dilations[0],
            norm_layer=norm_layer,
        )
        self.block5 = self._make_layer(
            InvertedResidual,
            self.planes,
            inverted_residual_setting[5:],
            dilations[1],
            norm_layer=norm_layer,
        )
        self.last_inp_channels = self.planes

        # building last several layers
        # features = list()
        # features.append(_ConvBNReLU(input_channels, last_channels, 1, relu6=True, norm_layer=norm_layer))
        # features.append(nn.AdaptiveAvgPool2d(1))
        # self.features = nn.Sequential(*features)
        #
        # self.classifier = nn.Sequential(
        #     nn.Dropout2d(0.2),
        #     nn.Linear(last_channels, num_classes))

        # weight initialization
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out")
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def _make_layer(
        self,
        block,
        planes,
        inverted_residual_setting,
        dilation=1,
        norm_layer=nn.BatchNorm2d,
    ):
        features = list()
        for t, c, n, s in inverted_residual_setting:
            out_channels = int(c * self.multiplier)
            stride = s if dilation == 1 else 1
            features.append(
                block(planes, out_channels, stride, t, dilation, norm_layer)
            )
            planes = out_channels
            for i in range(n - 1):
                features.append(
                    block(planes, out_channels, 1, t, norm_layer=norm_layer)
                )
                planes = out_channels
        self.planes = planes
        return nn.Sequential(*features)

    def forward(self, x, side_feature):
        x = self.conv1(x)
        x = x + side_feature
        x = self.block1(x)
        c1 = self.block2(x)
        c2 = self.block3(c1)
        c3 = self.block4(c2)
        c4 = self.block5(c3)
        # x = self.features(x)
        # x = self.classifier(x.view(x.size(0), x.size(1)))
        return c1, c2, c3, c4


def mobilenet_v2(norm_layer=nn.BatchNorm2d):
    return MobileNetV2(norm_layer=norm_layer)


# ============  Segmentor  ============


class LRASPP(nn.Module):
    """Lite R-ASPP"""

    def __init__(self, in_channels, out_channels, norm_layer=nn.BatchNorm2d, **kwargs):
        super(LRASPP, self).__init__()
        self.b0 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            norm_layer(out_channels),
            nn.ReLU(True),
        )
        self.b1 = nn.Sequential(
            nn.AdaptiveAvgPool2d((2, 2)),
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.Sigmoid(),
        )

    def forward(self, x):
        size = x.size()[2:]
        feat1 = self.b0(x)
        feat2 = self.b1(x)
        feat2 = F.interpolate(feat2, size, mode="bilinear", align_corners=True)
        x = feat1 * feat2
        return x


class MobileSeg(nn.Module):
    def __init__(self, nclass=1, **kwargs):
        super(MobileSeg, self).__init__()
        self.backbone = mobilenet_v2()
        self.lraspp = LRASPP(320, 128)
        self.fusion_conv1 = nn.Conv2d(128, 16, 1, 1, 0)
        self.fusion_conv2 = nn.Conv2d(24, 16, 1, 1, 0)
        self.head = nn.Conv2d(16, nclass, 1, 1, 0)
        self.aux_head = nn.Conv2d(16, nclass, 1, 1, 0)

    def forward(self, x, side_feature):
        x4, _, _, x8 = self.backbone(x, side_feature)
        x8 = self.lraspp(x8)
        x8 = F.interpolate(x8, x4.size()[2:], mode="bilinear", align_corners=True)
        x8 = self.fusion_conv1(x8)
        pred_aux = self.aux_head(x8)

        x4 = self.fusion_conv2(x4)
        x = x4 + x8
        pred = self.head(x)
        return pred, pred_aux, x

    def load_pretrained_weights(self, path_to_weights=" "):
        backbone_state_dict = self.backbone.state_dict()
        pretrained_state_dict = torch.load(path_to_weights, map_location="cpu")
        ckpt_keys = set(pretrained_state_dict.keys())
        own_keys = set(backbone_state_dict.keys())
        missing_keys = own_keys - ckpt_keys
        unexpected_keys = ckpt_keys - own_keys
        print("Loading Mobilnet V2")
        print("Missing Keys: ", missing_keys)
        print("Unexpected Keys: ", unexpected_keys)
        backbone_state_dict.update(pretrained_state_dict)
        self.backbone.load_state_dict(backbone_state_dict, strict=False)


class ScaleLayer(nn.Module):
    def __init__(self, init_value=1.0, lr_mult=1):
        super().__init__()
        self.lr_mult = lr_mult
        self.scale = nn.Parameter(
            torch.full((1,), init_value / lr_mult, dtype=torch.float32)
        )

    def forward(self, x):
        scale = torch.abs(self.scale * self.lr_mult)
        return x * scale


# ============ Interactive Segmentor  ============


class BaselineModel(nn.Module):
    def __init__(self, backbone_lr_mult=0.1, norm_layer=nn.BatchNorm2d, **kwargs):
        super().__init__()
        self.feature_extractor = MobileSeg()
        side_feature_ch = 32
        mt_layers = [
            nn.Conv2d(
                in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1
            ),
            nn.LeakyReLU(negative_slope=0.2),
            nn.Conv2d(
                in_channels=16,
                out_channels=side_feature_ch,
                kernel_size=3,
                stride=1,
                padding=1,
            ),
            ScaleLayer(init_value=0.05, lr_mult=1),
        ]
        self.maps_transform = nn.Sequential(*mt_layers)

    def backbone_forward(self, image, coord_features=None):
        mask, mask_aux, feature = self.feature_extractor(image, coord_features)
        return {"instances": mask, "instances_aux": mask_aux, "feature": feature}

    def prepare_input(self, image):
        prev_mask = torch.zeros_like(image)[:, :1, :, :]
        return image, prev_mask

    def forward(self, image, coarse_mask):
        image, prev_mask = self.prepare_input(image)
        coord_features = torch.cat((prev_mask, coarse_mask, coarse_mask * 0.0), dim=1)
        click_map = coord_features[:, 1:, :, :]

        coord_features = self.maps_transform(coord_features)
        outputs = self.backbone_forward(image, coord_features)

        pred = nn.functional.interpolate(
            outputs["instances"],
            size=image.size()[2:],
            mode="bilinear",
            align_corners=True,
        )

        outputs["instances"] = torch.sigmoid(pred)
        return outputs
