import numpy as np
import paddle
import paddle.nn as nn
from libs.tools import change_default_args
from libs.tools import GroupNorm
from libs.nn import Empty

class RPN(nn.Layer):
    def __init__(self,
                 use_norm=True,
                 num_class=2,
                 layer_nums=[3, 5, 5],
                 layer_strides=[2, 2, 2],
                 num_filters=[128, 128, 256],
                 upsample_strides=[1, 2, 4],
                 num_upsample_filters=[256, 256, 256],
                 num_input_filters=128,
                 num_anchor_per_loc=2,
                 encode_background_as_zeros=True,
                 use_direction_classifier=True,
                 use_groupnorm=False,
                 num_groups=32,
                 use_bev=False,
                 box_code_size=7,
                 name='rpn'):
        super(RPN, self).__init__()
        self._num_anchor_per_loc = num_anchor_per_loc
        self._use_direction_classifier = use_direction_classifier
        self._use_bev = use_bev
        assert len(layer_nums) == 3
        assert len(layer_strides) == len(layer_nums)
        assert len(num_filters) == len(layer_nums)
        assert len(upsample_strides) == len(layer_nums)
        assert len(num_upsample_filters) == len(layer_nums)
        factors = []
        for i in range(len(layer_nums)):
            assert int(np.prod(layer_strides[:i + 1])) % upsample_strides[i] == 0
            factors.append(np.prod(layer_strides[:i + 1]) // upsample_strides[i])
        assert all([x == factors[0] for x in factors])
        if use_norm:
            if use_groupnorm:
                BatchNorm2d = change_default_args(
                    num_groups=num_groups, epsilon=1e-3)(GroupNorm)
            else:
                BatchNorm2d = change_default_args(
                    epsilon=1e-3, momentum=0.01)(nn.BatchNorm2D)
            Conv2d = change_default_args(bias_attr=False)(nn.Conv2D)
            ConvTranspose2d = change_default_args(bias_attr=False)(
                nn.Conv2DTranspose)
        else:
            BatchNorm2d = Empty
            Conv2d = change_default_args(bias_attr=True)(nn.Conv2D)
            ConvTranspose2d = change_default_args(bias_attr=True)(
                nn.Conv2DTranspose)

        # note that when stride > 1, conv2d with same padding isn't
        # equal to pad-conv2d. we should use pad-conv2d.
        block2_input_filters = num_filters[0]
        if use_bev:
            self.bev_extractor = nn.Sequential(
                Conv2d(6, 32, 3, padding=1),
                BatchNorm2d(32),
                nn.ReLU(),
                # nn.MaxPool2d(2, 2),
                Conv2d(32, 64, 3, padding=1),
                BatchNorm2d(64),
                nn.ReLU(),
                nn.MaxPool2D(2, 2),
            )
            block2_input_filters += 64
        block1_layer = [
            nn.Pad2D(1),
            Conv2d(num_input_filters, num_filters[0], 3, stride=layer_strides[0]),
            BatchNorm2d(num_filters[0]),
            nn.ReLU(),
        ]
        for i in range(layer_nums[0]):
            block1_layer.append(
                Conv2d(num_filters[0], num_filters[0], 3, padding=1))
            block1_layer.append(BatchNorm2d(num_filters[0]))
            block1_layer.append(nn.ReLU())
        self.block1 = nn.Sequential(*block1_layer)
        self.deconv1 = nn.Sequential(
            ConvTranspose2d(
                num_filters[0],
                num_upsample_filters[0],
                upsample_strides[0],
                stride=upsample_strides[0]),
            BatchNorm2d(num_upsample_filters[0]),
            nn.ReLU(),
        )
        block2_layer = [
            nn.Pad2D(1),
            Conv2d(block2_input_filters, num_filters[1], 3, stride=layer_strides[1]),
            BatchNorm2d(num_filters[1]),
            nn.ReLU()
        ]
        for i in range(layer_nums[1]):
            block2_layer.append(
                Conv2d(num_filters[1], num_filters[1], 3, padding=1))
            block2_layer.append(BatchNorm2d(num_filters[1]))
            block2_layer.append(nn.ReLU())
        self.block2 = nn.Sequential(*block2_layer)
        self.deconv2 = nn.Sequential(
            ConvTranspose2d(
                num_filters[1],
                num_upsample_filters[1],
                upsample_strides[1],
                stride=upsample_strides[1]),
            BatchNorm2d(num_upsample_filters[1]),
            nn.ReLU(),
        )
        block3_layer = [
            nn.Pad2D(1),
            Conv2d(num_filters[1], num_filters[2], 3, stride=layer_strides[2]),
            BatchNorm2d(num_filters[2]),
            nn.ReLU(),
        ]
        for i in range(layer_nums[2]):
            block3_layer.append(
                Conv2d(num_filters[2], num_filters[2], 3, padding=1))
            block3_layer.append(BatchNorm2d(num_filters[2]))
            block3_layer.append(nn.ReLU())
        self.block3 = nn.Sequential(*block3_layer)
        self.deconv3 = nn.Sequential(
            ConvTranspose2d(
                num_filters[2],
                num_upsample_filters[2],
                upsample_strides[2],
                stride=upsample_strides[2]),
            BatchNorm2d(num_upsample_filters[2]),
            nn.ReLU(),
        )
        if encode_background_as_zeros:
            num_cls = num_anchor_per_loc * num_class
        else:
            num_cls = num_anchor_per_loc * (num_class + 1)
        self.conv_cls = nn.Conv2D(sum(num_upsample_filters), num_cls, 1)
        self.conv_box = nn.Conv2D(
            sum(num_upsample_filters), num_anchor_per_loc * box_code_size, 1)
        if use_direction_classifier:
            self.conv_dir_cls = nn.Conv2D(
                sum(num_upsample_filters), num_anchor_per_loc * 2, 1)

    def forward(self, x, bev=None):
        x = self.block1(x)
        up1 = self.deconv1(x)
        if self._use_bev:
            bev[:, -1] = paddle.clip(
                paddle.log(1 + bev[:, -1]) / np.log(16.0), max=1.0)
            x = paddle.concat([x, self.bev_extractor(bev)], axis=1)
        x = self.block2(x)
        up2 = self.deconv2(x)
        x = self.block3(x)
        up3 = self.deconv3(x)
        x = paddle.concat([up1, up2, up3], axis=1)
        box_preds = self.conv_box(x)
        cls_preds = self.conv_cls(x)
        # [N, C, y(H), x(W)]
        box_preds = box_preds.transpose((0, 2, 3, 1))
        cls_preds = cls_preds.transpose((0, 2, 3, 1))
        ret_dict = {
            "box_preds": box_preds,
            "cls_preds": cls_preds,
        }
        if self._use_direction_classifier:
            dir_cls_preds = self.conv_dir_cls(x)
            dir_cls_preds = dir_cls_preds.transpose((0, 2, 3, 1))
            ret_dict["dir_cls_preds"] = dir_cls_preds
        return ret_dict

if __name__ == '__main__':
    inputs = paddle.zeros([1, 128, 496, 432])
    rpn = RPN()
    result = RPN(inputs)
    print('succeed!')
