# -*- coding:utf-8 -*-
# This repo is licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import megengine.functional as F
import megengine.module as M

import layers


class RPN(M.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.box_coder = layers.BoxCoder(cfg.rpn_reg_mean, cfg.rpn_reg_std)

        # check anchor settings
        assert len(set(len(x) for x in cfg.anchor_scales)) == 1
        assert len(set(len(x) for x in cfg.anchor_ratios)) == 1
        self.num_cell_anchors = len(cfg.anchor_scales[0]) * len(cfg.anchor_ratios[0])

        rpn_channel = cfg.rpn_channel
        self.in_features = cfg.rpn_in_features

        self.anchor_generator = layers.AnchorBoxGenerator(
            anchor_scales=cfg.anchor_scales,
            anchor_ratios=cfg.anchor_ratios,
            strides=cfg.rpn_stride,
            offset=self.cfg.anchor_offset,
        )

        self.matcher = layers.Matcher(
            cfg.match_thresholds, cfg.match_labels, cfg.match_allow_low_quality
        )

        self.rpn_conv = M.Conv2d(256, rpn_channel, kernel_size=3, stride=1, padding=1)
        self.rpn_cls_score = M.Conv2d(
            rpn_channel, self.num_cell_anchors, kernel_size=1, stride=1
        )
        self.rpn_bbox_offsets = M.Conv2d(
            rpn_channel, self.num_cell_anchors * 4, kernel_size=1, stride=1
        )

        for l in [self.rpn_conv, self.rpn_cls_score, self.rpn_bbox_offsets]:
            M.init.normal_(l.weight, std=0.01)
            M.init.fill_(l.bias, 0)

    def forward(self, features, im_info, boxes=None):
        # prediction
        features = [features[x] for x in self.in_features]

        # get anchors
        anchors_list = self.anchor_generator(features)

        pred_cls_logit_list = []
        pred_bbox_offset_list = []
        for x in features:
            t = F.relu(self.rpn_conv(x))
            scores = self.rpn_cls_score(t)
            pred_cls_logit_list.append(
                scores.reshape(
                    scores.shape[0],
                    self.num_cell_anchors,
                    scores.shape[2],
                    scores.shape[3],
                )
            )
            bbox_offsets = self.rpn_bbox_offsets(t)
            pred_bbox_offset_list.append(
                bbox_offsets.reshape(
                    bbox_offsets.shape[0],
                    self.num_cell_anchors,
                    4,
                    bbox_offsets.shape[2],
                    bbox_offsets.shape[3],
                )
            )
        # get rois from the predictions
        rpn_rois = self.find_top_rpn_proposals(
            pred_cls_logit_list, pred_bbox_offset_list, anchors_list, im_info
        )

        if self.training:
            rpn_labels, rpn_offsets = self.get_ground_truth(
                anchors_list, boxes, im_info[:, 4].astype("int32")
            )
            pred_cls_logits, pred_bbox_offsets = self.merge_rpn_score_box(
                pred_cls_logit_list, pred_bbox_offset_list
            )

            fg_mask = rpn_labels > 0
            valid_mask = rpn_labels >= 0
            num_valid = valid_mask.sum()

            # rpn classification loss
            loss_rpn_cls = 
F.loss.binary_cross_entropy(
                pred_cls_logits[valid_mask], rpn_labels[