#!/usr/bin/env python
# -*- encoding: utf-8 -*-

"""
@Author  :   Peike Li
@Contact :   peike.li@yahoo.com
@File    :   AugmentCE2P.py
@Time    :   8/4/19 3:35 PM
@Desc    :
@License :   This source code is licensed under the license found in the
             LICENSE file in the root directory of this source tree.
"""

import torch
import torch.nn as nn

from torch.nn import BatchNorm2d, functional as F, LeakyReLU

affine_par = True
pretrained_settings = {
    "resnet101": {
        "imagenet": {
            "input_space": "BGR",
            "input_size": [3, 224, 224],
            "input_range": [0, 1],
            "mean": [0.406, 0.456, 0.485],
            "std": [0.225, 0.224, 0.229],
            "num_classes": 1000,
        }
    },
}


def conv3x3(in_planes, out_planes, stride=1):
    "3x3 convolution with padding"
    return nn.Conv2d(
        in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
    )


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(
        self,
        inplanes,
        planes,
        stride=1,
        dilation=1,
        downsample=None,
        fist_dilation=1,
        multi_grid=1,
    ):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(
            planes,
            planes,
            kernel_size=3,
            stride=stride,
            padding=dilation * multi_grid,
            dilation=dilation * multi_grid,
            bias=False,
        )
        self.bn2 = BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=False)
        self.relu_inplace = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.dilation = dilation
        self.stride = stride

    def forward(self, x):
        residual = x

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

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

        out = self.conv3(out)
        out = self.bn3(out)

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

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

        return out


class PSPModule(nn.Module):
    """
    Reference:
        Zhao, Hengshuang, et al. *"Pyramid scene parsing network."*
    """

    def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)):
        super(PSPModule, self).__init__()

        self.stages = []
        self.stages = nn.ModuleList(
            [self._make_stage(features, out_features, size) for size in sizes]
        )
        self.bottleneck = nn.Sequential(
            nn.Conv2d(
                features + len(sizes) * out_features,
                out_features,
                kernel_size=3,
                padding=1,
                dilation=1,
                bias=False,
            ),
            BatchNorm2d(out_features),
            LeakyReLU(),
        )

    def _make_stage(self, features, out_features, size):
        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
        conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False)
        return nn.Sequential(
            prior,
            conv,
            # bn
            BatchNorm2d(out_features),
            LeakyReLU(),
        )

    def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        priors = [
            F.interpolate(
                input=stage(feats), size=(h, w), mode="bilinear", align_corners=True
            )
            for stage in self.stages
        ] + [feats]
        bottle = self.bottleneck(torch.cat(priors, 1))
        return bottle


class ASPPModule(nn.Module):
    """
    Reference:
        Chen, Liang-Chieh, et al. *"Rethinking Atrous Convolution for Semantic Image Segmentation."*
    """

    def __init__(
        self, features, inner_features=256, out_features=512, dilations=(12, 24, 36)
    ):
        super(ASPPModule, self).__init__()

        self.conv1 = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Conv2d(
                features,
                inner_features,
                kernel_size=1,
                padding=0,
                dilation=1,
                bias=False,
            ),
            #    InPlaceABNSync(inner_features)
            BatchNorm2d(inner_features),
            LeakyReLU(),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(
                features,
                inner_features,
                kernel_size=1,
                padding=0,
                dilation=1,
                bias=False,
            ),
            BatchNorm2d(inner_features),
            LeakyReLU(),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(
                features,
                inner_features,
                kernel_size=3,
                padding=dilations[0],
                dilation=dilations[0],
                bias=False,
            ),
            BatchNorm2d(inner_features),
            LeakyReLU(),
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(
                features,
                inner_features,
                kernel_size=3,
                padding=dilations[1],
                dilation=dilations[1],
                bias=False,
            ),
            BatchNorm2d(inner_features),
            LeakyReLU(),
        )
        self.conv5 = nn.Sequential(
            nn.Conv2d(
                features,
                inner_features,
                kernel_size=3,
                padding=dilations[2],
                dilation=dilations[2],
                bias=False,
            ),
            BatchNorm2d(inner_features),
            LeakyReLU(),
        )

        self.bottleneck = nn.Sequential(
            nn.Conv2d(
                inner_features * 5,
                out_features,
                kernel_size=1,
                padding=0,
                dilation=1,
                bias=False,
            ),
            BatchNorm2d(inner_features),
            LeakyReLU(),
            nn.Dropout2d(0.1),
        )

    def forward(self, x):
        _, _, h, w = x.size()

        feat1 = F.interpolate(
            self.conv1(x), size=(h, w), mode="bilinear", align_corners=True
        )

        feat2 = self.conv2(x)
        feat3 = self.conv3(x)
        feat4 = self.conv4(x)
        feat5 = self.conv5(x)
        out = torch.cat((feat1, feat2, feat3, feat4, feat5), 1)

        bottle = self.bottleneck(out)
        return bottle


class Edge_Module(nn.Module):
    """
    Edge Learning Branch
    """

    def __init__(self, in_fea=[256, 512, 1024], mid_fea=256, out_fea=2):
        super(Edge_Module, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_fea[0], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False
            ),
            BatchNorm2d(mid_fea),
            LeakyReLU(),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(
                in_fea[1], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False
            ),
            BatchNorm2d(mid_fea),
            LeakyReLU(),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(
                in_fea[2], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False
            ),
            BatchNorm2d(mid_fea),
            LeakyReLU(),
        )
        self.conv4 = nn.Conv2d(
            mid_fea, out_fea, kernel_size=3, padding=1, dilation=1, bias=True
        )
        # self.conv5 = nn.Conv2d(out_fea * 3, out_fea, kernel_size=1, padding=0, dilation=1, bias=True)

    def forward(self, x1, x2, x3):
        _, _, h, w = x1.size()

        edge1_fea = self.conv1(x1)
        # edge1 = self.conv4(edge1_fea)
        edge2_fea = self.conv2(x2)
        edge2 = self.conv4(edge2_fea)
        edge3_fea = self.conv3(x3)
        edge3 = self.conv4(edge3_fea)

        edge2_fea = F.interpolate(
            edge2_fea, size=(h, w), mode="bilinear", align_corners=True
        )
        edge3_fea = F.interpolate(
            edge3_fea, size=(h, w), mode="bilinear", align_corners=True
        )
        edge2 = F.interpolate(edge2, size=(h, w), mode="bilinear", align_corners=True)
        edge3 = F.interpolate(edge3, size=(h, w), mode="bilinear", align_corners=True)

        # edge = torch.cat([edge1, edge2, edge3], dim=1)
        edge_fea = torch.cat([edge1_fea, edge2_fea, edge3_fea], dim=1)
        # edge = self.conv5(edge)

        # return edge, edge_fea
        return edge_fea


class Decoder_Module(nn.Module):
    """
    Parsing Branch Decoder Module.
    """

    def __init__(self, num_classes):
        super(Decoder_Module, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(512, 256, kernel_size=1, padding=0, dilation=1, bias=False),
            BatchNorm2d(256),
            LeakyReLU(),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(
                256, 48, kernel_size=1, stride=1, padding=0, dilation=1, bias=False
            ),
            BatchNorm2d(48),
            LeakyReLU(),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(304, 256, kernel_size=1, padding=0, dilation=1, bias=False),
            BatchNorm2d(256),
            LeakyReLU(),
            nn.Conv2d(256, 256, kernel_size=1, padding=0, dilation=1, bias=False),
            BatchNorm2d(256),
            LeakyReLU(),
        )

        # self.conv4 = nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)

    def forward(self, xt, xl):
        _, _, h, w = xl.size()
        xt = F.interpolate(
            self.conv1(xt), size=(h, w), mode="bilinear", align_corners=True
        )
        xl = self.conv2(xl)
        x = torch.cat([xt, xl], dim=1)
        x = self.conv3(x)
        # seg = self.conv4(x)
        # return seg, x
        return x


class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes):
        self.inplanes = 128
        super(ResNet, self).__init__()
        self.conv1 = conv3x3(3, 64, stride=2)
        self.bn1 = BatchNorm2d(64)
        self.relu1 = nn.ReLU(inplace=False)
        self.conv2 = conv3x3(64, 64)
        self.bn2 = BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=False)
        self.conv3 = conv3x3(64, 128)
        self.bn3 = BatchNorm2d(128)
        self.relu3 = nn.ReLU(inplace=False)

        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(
            block, 512, layers[3], stride=1, dilation=2, multi_grid=(1, 1, 1)
        )

        self.context_encoding = PSPModule(2048, 512)

        self.edge = Edge_Module()
        self.decoder = Decoder_Module(num_classes)

        self.fushion = nn.Sequential(
            nn.Conv2d(1024, 256, kernel_size=1, padding=0, dilation=1, bias=False),
            BatchNorm2d(256),
            LeakyReLU(),
            nn.Dropout2d(0.1),
            nn.Conv2d(
                256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True
            ),
        )

    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, multi_grid=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(
                    self.inplanes,
                    planes * block.expansion,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                BatchNorm2d(planes * block.expansion, affine=affine_par),
            )

        layers = []
        generate_multi_grid = lambda index, grids: (
            grids[index % len(grids)] if isinstance(grids, tuple) else 1
        )
        layers.append(
            block(
                self.inplanes,
                planes,
                stride,
                dilation=dilation,
                downsample=downsample,
                multi_grid=generate_multi_grid(0, multi_grid),
            )
        )
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(
                block(
                    self.inplanes,
                    planes,
                    dilation=dilation,
                    multi_grid=generate_multi_grid(i, multi_grid),
                )
            )

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.relu1(self.bn1(self.conv1(x)))
        x = self.relu2(self.bn2(self.conv2(x)))
        x = self.relu3(self.bn3(self.conv3(x)))
        x = self.maxpool(x)
        x2 = self.layer1(x)
        x3 = self.layer2(x2)
        x4 = self.layer3(x3)
        x5 = self.layer4(x4)
        x = self.context_encoding(x5)
        # parsing_result, parsing_fea = self.decoder(x, x2)
        parsing_fea = self.decoder(x, x2)
        # Edge Branch
        # edge_result, edge_fea = self.edge(x2, x3, x4)
        edge_fea = self.edge(x2, x3, x4)
        # Fusion Branch
        x = torch.cat([parsing_fea, edge_fea], dim=1)
        fusion_result = self.fushion(x)
        # return [[parsing_result, fusion_result], [edge_result]]
        return fusion_result


def initialize_pretrained_model(
    model, settings, pretrained="./models/resnet101-imagenet.pth"
):
    model.input_space = settings["input_space"]
    model.input_size = settings["input_size"]
    model.input_range = settings["input_range"]
    model.mean = settings["mean"]
    model.std = settings["std"]

    if pretrained is not None:
        saved_state_dict = torch.load(pretrained)
        new_params = model.state_dict().copy()
        for i in saved_state_dict:
            i_parts = i.split(".")
            if not i_parts[0] == "fc":
                new_params[".".join(i_parts[0:])] = saved_state_dict[i]
        model.load_state_dict(new_params)


def resnet101(num_classes=20, pretrained="./models/resnet101-imagenet.pth"):
    model = ResNet(Bottleneck, [3, 4, 23, 3], num_classes)
    settings = pretrained_settings["resnet101"]["imagenet"]
    initialize_pretrained_model(model, settings, pretrained)
    return model