# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# The gca code was heavily based on https://github.com/Yaoyi-Li/GCA-Matting
# and https://github.com/open-mmlab/mmediting

import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from paddleseg.cvlibs import manager, param_init
from paddleseg.utils import utils

from ppmatting.models.layers import GuidedCxtAtten


class ResNet_D(nn.Layer):
    def __init__(self,
                 input_channels,
                 layers,
                 late_downsample=False,
                 pretrained=None):

        super().__init__()

        self.pretrained = pretrained

        self._norm_layer = nn.BatchNorm
        self.inplanes = 64
        self.late_downsample = late_downsample
        self.midplanes = 64 if late_downsample else 32
        self.start_stride = [1, 2, 1, 2] if late_downsample else [2, 1, 2, 1]
        self.conv1 = nn.utils.spectral_norm(
            nn.Conv2D(
                input_channels,
                32,
                kernel_size=3,
                stride=self.start_stride[0],
                padding=1,
                bias_attr=False))
        self.conv2 = nn.utils.spectral_norm(
            nn.Conv2D(
                32,
                self.midplanes,
                kernel_size=3,
                stride=self.start_stride[1],
                padding=1,
                bias_attr=False))
        self.conv3 = nn.utils.spectral_norm(
            nn.Conv2D(
                self.midplanes,
                self.inplanes,
                kernel_size=3,
                stride=self.start_stride[2],
                padding=1,
                bias_attr=False))
        self.bn1 = self._norm_layer(32)
        self.bn2 = self._norm_layer(self.midplanes)
        self.bn3 = self._norm_layer(self.inplanes)
        self.activation = nn.ReLU()
        self.layer1 = self._make_layer(
            BasicBlock, 64, layers[0], stride=self.start_stride[3])
        self.layer2 = self._make_layer(BasicBlock, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(BasicBlock, 256, layers[2], stride=2)
        self.layer_bottleneck = self._make_layer(
            BasicBlock, 512, layers[3], stride=2)

        self.init_weight()

    def _make_layer(self, block, planes, block_num, stride=1):
        if block_num == 0:
            return nn.Sequential(nn.Identity())
        norm_layer = self._norm_layer
        downsample = None
        if stride != 1:
            downsample = nn.Sequential(
                nn.AvgPool2D(2, stride),
                nn.utils.spectral_norm(
                    conv1x1(self.inplanes, planes * block.expansion)),
                norm_layer(planes * block.expansion), )
        elif self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.utils.spectral_norm(
                    conv1x1(self.inplanes, planes * block.expansion, stride)),
                norm_layer(planes * block.expansion), )

        layers = [block(self.inplanes, planes, stride, downsample, norm_layer)]
        self.inplanes = planes * block.expansion
        for _ in range(1, block_num):
            layers.append(block(self.inplanes, planes, norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.activation(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x1 = self.activation(x)  # N x 32 x 256 x 256
        x = self.conv3(x1)
        x = self.bn3(x)
        x2 = self.activation(x)  # N x 64 x 128 x 128

        x3 = self.layer1(x2)  # N x 64 x 128 x 128
        x4 = self.layer2(x3)  # N x 128 x 64 x 64
        x5 = self.layer3(x4)  # N x 256 x 32 x 32
        x = self.layer_bottleneck(x5)  # N x 512 x 16 x 16

        return x, (x1, x2, x3, x4, x5)

    def init_weight(self):

        for layer in self.sublayers():
            if isinstance(layer, nn.Conv2D):

                if hasattr(layer, "weight_orig"):
                    param = layer.weight_orig
                else:
                    param = layer.weight
                param_init.xavier_uniform(param)

            elif isinstance(layer, (nn.BatchNorm, nn.SyncBatchNorm)):
                param_init.constant_init(layer.weight, value=1.0)
                param_init.constant_init(layer.bias, value=0.0)

            elif isinstance(layer, BasicBlock):
                param_init.constant_init(layer.bn2.weight, value=0.0)

        if self.pretrained is not None:
            utils.load_pretrained_model(self, self.pretrained)


@manager.MODELS.add_component
class ResShortCut_D(ResNet_D):
    def __init__(self,
                 input_channels,
                 layers,
                 late_downsample=False,
                 pretrained=None):
        super().__init__(
            input_channels,
            layers,
            late_downsample=late_downsample,
            pretrained=pretrained)

        self.shortcut_inplane = [input_channels, self.midplanes, 64, 128, 256]
        self.shortcut_plane = [32, self.midplanes, 64, 128, 256]

        self.shortcut = nn.LayerList()
        for stage, inplane in enumerate(self.shortcut_inplane):
            self.shortcut.append(
                self._make_shortcut(inplane, self.shortcut_plane[stage]))

    def _make_shortcut(self, inplane, planes):
        return nn.Sequential(
            nn.utils.spectral_norm(
                nn.Conv2D(
                    inplane, planes, kernel_size=3, padding=1,
                    bias_attr=False)),
            nn.ReLU(),
            self._norm_layer(planes),
            nn.utils.spectral_norm(
                nn.Conv2D(
                    planes, planes, kernel_size=3, padding=1, bias_attr=False)),
            nn.ReLU(),
            self._norm_layer(planes))

    def forward(self, x):

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.activation(out)
        out = self.conv2(out)
        out = self.bn2(out)
        x1 = self.activation(out)  # N x 32 x 256 x 256
        out = self.conv3(x1)
        out = self.bn3(out)
        out = self.activation(out)

        x2 = self.layer1(out)  # N x 64 x 128 x 128
        x3 = self.layer2(x2)  # N x 128 x 64 x 64
        x4 = self.layer3(x3)  # N x 256 x 32 x 32
        out = self.layer_bottleneck(x4)  # N x 512 x 16 x 16

        fea1 = self.shortcut[0](x)  # input image and trimap
        fea2 = self.shortcut[1](x1)
        fea3 = self.shortcut[2](x2)
        fea4 = self.shortcut[3](x3)
        fea5 = self.shortcut[4](x4)

        return out, {
            'shortcut': (fea1, fea2, fea3, fea4, fea5),
            'image': x[:, :3, ...]
        }


@manager.MODELS.add_component
class ResGuidedCxtAtten(ResNet_D):
    def __init__(self,
                 input_channels,
                 layers,
                 late_downsample=False,
                 pretrained=None):
        super().__init__(
            input_channels,
            layers,
            late_downsample=late_downsample,
            pretrained=pretrained)
        self.input_channels = input_channels
        self.shortcut_inplane = [input_channels, self.midplanes, 64, 128, 256]
        self.shortcut_plane = [32, self.midplanes, 64, 128, 256]

        self.shortcut = nn.LayerList()
        for stage, inplane in enumerate(self.shortcut_inplane):
            self.shortcut.append(
                self._make_shortcut(inplane, self.shortcut_plane[stage]))

        self.guidance_head = nn.Sequential(
            nn.Pad2D(
                1, mode="reflect"),
            nn.utils.spectral_norm(
                nn.Conv2D(
                    3, 16, kernel_size=3, padding=0, stride=2,
                    bias_attr=False)),
            nn.ReLU(),
            self._norm_layer(16),
            nn.Pad2D(
                1, mode="reflect"),
            nn.utils.spectral_norm(
                nn.Conv2D(
                    16, 32, kernel_size=3, padding=0, stride=2,
                    bias_attr=False)),
            nn.ReLU(),
            self._norm_layer(32),
            nn.Pad2D(
                1, mode="reflect"),
            nn.utils.spectral_norm(
                nn.Conv2D(
                    32,
                    128,
                    kernel_size=3,
                    padding=0,
                    stride=2,
                    bias_attr=False)),
            nn.ReLU(),
            self._norm_layer(128))

        self.gca = GuidedCxtAtten(128, 128)

        self.init_weight()

    def init_weight(self):

        for layer in self.sublayers():
            if isinstance(layer, nn.Conv2D):
                initializer = nn.initializer.XavierUniform()
                if hasattr(layer, "weight_orig"):
                    param = layer.weight_orig
                else:
                    param = layer.weight
                initializer(param, param.block)

            elif isinstance(layer, (nn.BatchNorm, nn.SyncBatchNorm)):
                param_init.constant_init(layer.weight, value=1.0)
                param_init.constant_init(layer.bias, value=0.0)

            elif isinstance(layer, BasicBlock):
                param_init.constant_init(layer.bn2.weight, value=0.0)

        if self.pretrained is not None:
            utils.load_pretrained_model(self, self.pretrained)

    def _make_shortcut(self, inplane, planes):
        return nn.Sequential(
            nn.utils.spectral_norm(
                nn.Conv2D(
                    inplane, planes, kernel_size=3, padding=1,
                    bias_attr=False)),
            nn.ReLU(),
            self._norm_layer(planes),
            nn.utils.spectral_norm(
                nn.Conv2D(
                    planes, planes, kernel_size=3, padding=1, bias_attr=False)),
            nn.ReLU(),
            self._norm_layer(planes))

    def forward(self, x):

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.activation(out)
        out = self.conv2(out)
        out = self.bn2(out)
        x1 = self.activation(out)  # N x 32 x 256 x 256
        out = self.conv3(x1)
        out = self.bn3(out)
        out = self.activation(out)

        im_fea = self.guidance_head(
            x[:, :3, ...])  # downsample origin image and extract features
        if self.input_channels == 6:
            unknown = F.interpolate(
                x[:, 4:5, ...], scale_factor=1 / 8, mode='nearest')
        else:
            unknown = x[:, 3:, ...].equal(paddle.to_tensor([1.]))
            unknown = paddle.cast(unknown, dtype='float32')
            unknown = F.interpolate(unknown, scale_factor=1 / 8, mode='nearest')

        x2 = self.layer1(out)  # N x 64 x 128 x 128
        x3 = self.layer2(x2)  # N x 128 x 64 x 64
        x3 = self.gca(im_fea, x3, unknown)  # contextual attention
        x4 = self.layer3(x3)  # N x 256 x 32 x 32
        out = self.layer_bottleneck(x4)  # N x 512 x 16 x 16

        fea1 = self.shortcut[0](x)  # input image and trimap
        fea2 = self.shortcut[1](x1)
        fea3 = self.shortcut[2](x2)
        fea4 = self.shortcut[3](x3)
        fea5 = self.shortcut[4](x4)

        return out, {
            'shortcut': (fea1, fea2, fea3, fea4, fea5),
            'image_fea': im_fea,
            'unknown': unknown,
        }


class BasicBlock(nn.Layer):
    expansion = 1

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 downsample=None,
                 norm_layer=None):
        super().__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        self.conv1 = nn.utils.spectral_norm(conv3x3(inplanes, planes, stride))
        self.bn1 = norm_layer(planes)
        self.activation = nn.ReLU()
        self.conv2 = nn.utils.spectral_norm(conv3x3(planes, planes))
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

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

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

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

        out += identity
        out = self.activation(out)

        return out


def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2D(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=dilation,
        groups=groups,
        bias_attr=False,
        dilation=dilation)


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2D(
        in_planes, out_planes, kernel_size=1, stride=stride, bias_attr=False)