# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
#
# 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.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
import paddle
from paddle import nn
import paddle.nn.functional as F
from paddle import ParamAttr
from ppocr.modeling.backbones.det_mobilenet_v3 import ConvBNLayer


def get_bias_attr(k):
    stdv = 1.0 / math.sqrt(k * 1.0)
    initializer = paddle.nn.initializer.Uniform(-stdv, stdv)
    bias_attr = ParamAttr(initializer=initializer)
    return bias_attr


class Head(nn.Layer):
    def __init__(self, in_channels, kernel_list=[3, 2, 2], fix_nan=False, **kwargs):
        super(Head, self).__init__()

        self.conv1 = nn.Conv2D(
            in_channels=in_channels,
            out_channels=in_channels // 4,
            kernel_size=kernel_list[0],
            padding=int(kernel_list[0] // 2),
            weight_attr=ParamAttr(),
            bias_attr=False,
        )
        self.conv_bn1 = nn.BatchNorm(
            num_channels=in_channels // 4,
            param_attr=ParamAttr(initializer=paddle.nn.initializer.Constant(value=1.0)),
            bias_attr=ParamAttr(initializer=paddle.nn.initializer.Constant(value=1e-4)),
            act="relu",
        )

        self.conv2 = nn.Conv2DTranspose(
            in_channels=in_channels // 4,
            out_channels=in_channels // 4,
            kernel_size=kernel_list[1],
            stride=2,
            weight_attr=ParamAttr(initializer=paddle.nn.initializer.KaimingUniform()),
            bias_attr=get_bias_attr(in_channels // 4),
        )
        self.conv_bn2 = nn.BatchNorm(
            num_channels=in_channels // 4,
            param_attr=ParamAttr(initializer=paddle.nn.initializer.Constant(value=1.0)),
            bias_attr=ParamAttr(initializer=paddle.nn.initializer.Constant(value=1e-4)),
            act="relu",
        )
        self.conv3 = nn.Conv2DTranspose(
            in_channels=in_channels // 4,
            out_channels=1,
            kernel_size=kernel_list[2],
            stride=2,
            weight_attr=ParamAttr(initializer=paddle.nn.initializer.KaimingUniform()),
            bias_attr=get_bias_attr(in_channels // 4),
        )

        self.fix_nan = fix_nan
        self.is_repped = False

    def forward(self, x, return_f=False):
        x = self.conv1(x)
        x = self.conv_bn1(x)
        if self.fix_nan and self.training:
            x = paddle.where(paddle.isnan(x), paddle.zeros_like(x), x)
        x = self.conv2(x)
        x = self.conv_bn2(x)
        if self.fix_nan and self.training:
            x = paddle.where(paddle.isnan(x), paddle.zeros_like(x), x)
        if return_f is True:
            f = x
        x = self.conv3(x)
        x = F.sigmoid(x)
        if return_f is True:
            return x, f
        return x

    @paddle.no_grad()
    def rep(self):
        """Fuse Conv+BN and ConvTranspose+BN pairs for deployment."""
        if self.is_repped:
            return

        # conv1 (Conv2D, no bias) + conv_bn1 (BatchNorm, act=relu)
        self.conv1 = self._fuse_conv_bn(self.conv1, self.conv_bn1)
        self.conv_bn1 = nn.ReLU()

        # conv2 (Conv2DTranspose, has bias) + conv_bn2 (BatchNorm, act=relu)
        self.conv2 = self._fuse_convtranspose_bn(self.conv2, self.conv_bn2)
        self.conv_bn2 = nn.ReLU()

        self.is_repped = True

    @staticmethod
    @paddle.no_grad()
    def _fuse_conv_bn(conv, bn):
        """Fuse Conv2D + BatchNorm into Conv2D with bias."""
        gamma = bn.weight
        std = paddle.sqrt(bn._variance + bn._epsilon)
        scale = gamma / std

        w = conv.weight * scale[:, None, None, None]
        b = bn.bias - bn._mean * scale

        fused = nn.Conv2D(
            conv._in_channels,
            conv._out_channels,
            conv._kernel_size,
            stride=conv._stride,
            padding=conv._padding,
            dilation=conv._dilation,
            groups=conv._groups,
        )
        fused.weight.set_value(w)
        fused.bias.set_value(b)
        return fused

    @staticmethod
    @paddle.no_grad()
    def _fuse_convtranspose_bn(conv, bn):
        """Fuse Conv2DTranspose + BatchNorm into Conv2DTranspose with bias.

        Conv2DTranspose weight shape: [in_ch, out_ch/groups, kH, kW]
        BN scale applies on out_ch, i.e. axis=1.
        """
        gamma = bn.weight
        std = paddle.sqrt(bn._variance + bn._epsilon)
        scale = gamma / std

        # axis=1 for ConvTranspose (output channel dimension)
        w = conv.weight * scale[None, :, None, None]
        b = bn.bias - bn._mean * scale
        if conv.bias is not None:
            b = b + conv.bias * scale

        fused = nn.Conv2DTranspose(
            conv._in_channels,
            conv._out_channels,
            conv._kernel_size,
            stride=conv._stride,
            padding=conv._padding,
            dilation=conv._dilation,
            groups=conv._groups,
        )
        fused.weight.set_value(w)
        fused.bias.set_value(b)
        return fused


class DBHead(nn.Layer):
    """
    Differentiable Binarization (DB) for text detection:
        see https://arxiv.org/abs/1911.08947
    args:
        params(dict): super parameters for build DB network
    """

    def __init__(self, in_channels, k=50, aux_in_channels=0, **kwargs):
        super(DBHead, self).__init__()
        self.k = k
        self.is_repped = False
        self.binarize = Head(in_channels, **kwargs)
        self.thresh = Head(in_channels, **kwargs)
        self.aux_in_channels = aux_in_channels

        if aux_in_channels > 0:
            self._aux_upsample_scale = {
                "aux_p4": 4,  # 1/16 -> 1/4
                "aux_p3": 2,  # 1/8  -> 1/4
                "aux_p2": 1,  # 1/4  -> 1/4 (no-op)
            }
            # Create independent binarize + thresh Head pairs for each scale
            self.aux_binarize_p4 = Head(aux_in_channels, **kwargs)
            self.aux_thresh_p4 = Head(aux_in_channels, **kwargs)
            self.aux_binarize_p3 = Head(aux_in_channels, **kwargs)
            self.aux_thresh_p3 = Head(aux_in_channels, **kwargs)
            self.aux_binarize_p2 = Head(aux_in_channels, **kwargs)
            self.aux_thresh_p2 = Head(aux_in_channels, **kwargs)

    def step_function(self, x, y):
        return paddle.reciprocal(1 + paddle.exp(-self.k * (x - y)))

    def forward(self, x, targets=None):
        # Compatible with neck returning dict (training) or tensor (inference)
        if isinstance(x, dict):
            fuse = x["fuse"]
            aux_feats = {k: x[k] for k in ("aux_p4", "aux_p3", "aux_p2") if k in x}
        else:
            fuse = x
            aux_feats = {}

        shrink_maps = self.binarize(fuse)
        if not self.training:
            return {"maps": shrink_maps}

        threshold_maps = self.thresh(fuse)
        binary_maps = self.step_function(shrink_maps, threshold_maps)
        y = paddle.concat([shrink_maps, threshold_maps, binary_maps], axis=1)
        result = {"maps": y}

        if self.aux_in_channels > 0 and aux_feats:
            for key, feat in aux_feats.items():
                scale = self._aux_upsample_scale[key]
                if scale > 1:
                    feat = F.interpolate(
                        feat, scale_factor=scale, mode="bilinear", align_corners=False
                    )
                level = key[4:]  # 'p4', 'p3', 'p2'
                aux_binarize = getattr(self, "aux_binarize_" + level)
                aux_thresh_head = getattr(self, "aux_thresh_" + level)
                aux_shrink = aux_binarize(feat)
                aux_thresh = aux_thresh_head(feat)
                aux_binary = self.step_function(aux_shrink, aux_thresh)
                result["aux_maps_" + level] = paddle.concat(
                    [aux_shrink, aux_thresh, aux_binary], axis=1
                )

        return result

    def rep(self):
        """Fuse reparam structures in all sub-modules for deployment."""
        if self.is_repped:
            return
        for layer in self.sublayers():
            if isinstance(layer, Head):
                layer.rep()
        self.is_repped = True


class LocalModule(nn.Layer):
    def __init__(self, in_c, mid_c, use_distance=True):
        super(self.__class__, self).__init__()
        self.last_3 = ConvBNLayer(in_c + 1, mid_c, 3, 1, 1, act="relu")
        self.last_1 = nn.Conv2D(mid_c, 1, 1, 1, 0)

    def forward(self, x, init_map, distance_map):
        outf = paddle.concat([init_map, x], axis=1)
        # last Conv
        out = self.last_1(self.last_3(outf))
        return out


class PFHeadLocal(DBHead):
    def __init__(self, in_channels, k=50, mode="small", **kwargs):
        super(PFHeadLocal, self).__init__(in_channels, k, **kwargs)
        self.mode = mode

        self.up_conv = nn.Upsample(scale_factor=2, mode="nearest", align_mode=1)
        if self.mode == "large":
            self.cbn_layer = LocalModule(in_channels // 4, in_channels // 4)
        elif self.mode == "small":
            self.cbn_layer = LocalModule(in_channels // 4, in_channels // 8)

    def forward(self, x, targets=None):
        shrink_maps, f = self.binarize(x, return_f=True)
        base_maps = shrink_maps
        cbn_maps = self.cbn_layer(self.up_conv(f), shrink_maps, None)
        cbn_maps = F.sigmoid(cbn_maps)
        if not self.training:
            return {"maps": 0.5 * (base_maps + cbn_maps)}

        threshold_maps = self.thresh(x)
        binary_maps = self.step_function(shrink_maps, threshold_maps)
        y = paddle.concat([cbn_maps, threshold_maps, binary_maps], axis=1)
        return {"maps": y, "distance_maps": cbn_maps, "cbn_maps": binary_maps}
