"""
ViTDet (Vision Transformer for Object Detection).

Paper "Exploring Plain Vision Transformer Backbones for Object Detection",
https://arxiv.org/abs/2203.16527

Changes from original:
* Uses Faster R-CNN instead of Mask R-CNN, performing object detection without mask prediction
* Intended to use ViT-SAM (Vision Transformer with spatial modifications) as the backbone
"""

from typing import Any
from typing import Optional

from torch import nn
from torchvision.ops import MultiScaleRoIAlign
from torchvision.ops.feature_pyramid_network import LastLevelMaxPool

from birder.net.base import DetectorBackbone
from birder.net.detection.base import AnchorGenerator
from birder.net.detection.base import BackboneWithSimpleFPN
from birder.net.detection.faster_rcnn import Faster_RCNN
from birder.net.detection.faster_rcnn import FastRCNNConvFCHead
from birder.net.detection.faster_rcnn import FastRCNNPredictor
from birder.net.detection.faster_rcnn import RegionProposalNetwork
from birder.net.detection.faster_rcnn import RoIHeads
from birder.net.detection.faster_rcnn import RPNHead


class ViTDet(Faster_RCNN):
    default_size = (640, 640)
    auto_register = True

    # pylint: disable=too-many-locals
    def __init__(
        self,
        num_classes: int,
        backbone: DetectorBackbone,
        *,
        config: Optional[dict[str, Any]] = None,
        size: Optional[tuple[int, int]] = None,
        export_mode: bool = False,
    ) -> None:
        super().__init__(num_classes, backbone, config=config, size=size, export_mode=export_mode)
        assert self.config is None, "config not supported"

        fpn_width = 256
        rpn_pre_nms_top_n_train = 2000
        rpn_pre_nms_top_n_test = 1000
        rpn_post_nms_top_n_train = 2000
        rpn_post_nms_top_n_test = 1000
        rpn_fg_iou_thresh = 0.7
        rpn_bg_iou_thresh = 0.3
        rpn_batch_size_per_image = 256
        rpn_positive_fraction = 0.5
        rpn_nms_thresh = 0.7
        rpn_score_thresh = 0.0
        box_fg_iou_thresh = 0.5
        box_bg_iou_thresh = 0.5
        box_batch_size_per_image = 512
        box_positive_fraction = 0.25
        bbox_reg_weights = (10.0, 10.0, 5.0, 5.0)
        box_score_thresh = 0.05
        box_nms_thresh = 0.5
        box_detections_per_img = 100
        canonical_scale = 224

        self.backbone_with_fpn = BackboneWithSimpleFPN(self.backbone, fpn_width, extra_blocks=LastLevelMaxPool())

        anchor_sizes = [[32], [64], [128], [256], [512]]
        aspect_ratios = [[0.5, 1.0, 2.0]] * len(anchor_sizes)
        rpn_anchor_generator = AnchorGenerator(anchor_sizes, aspect_ratios)
        rpn_head = RPNHead(
            self.backbone_with_fpn.out_channels, rpn_anchor_generator.num_anchors_per_location()[0], conv_depth=2
        )

        rpn_pre_nms_top_n = {"training": rpn_pre_nms_top_n_train, "testing": rpn_pre_nms_top_n_test}
        rpn_post_nms_top_n = {"training": rpn_post_nms_top_n_train, "testing": rpn_post_nms_top_n_test}

        self.rpn = RegionProposalNetwork(
            rpn_anchor_generator,
            rpn_head,
            rpn_fg_iou_thresh,
            rpn_bg_iou_thresh,
            rpn_batch_size_per_image,
            rpn_positive_fraction,
            rpn_pre_nms_top_n,
            rpn_post_nms_top_n,
            rpn_nms_thresh,
            score_thresh=rpn_score_thresh,
        )
        box_roi_pool = MultiScaleRoIAlign(
            featmap_names=["stage1", "stage2", "stage3", "stage4"],  # Generated by the simple FPN
            output_size=canonical_scale // 32,
            sampling_ratio=2,
            canonical_scale=canonical_scale,
        )
        box_head = FastRCNNConvFCHead(
            (self.backbone_with_fpn.out_channels, canonical_scale // 32, canonical_scale // 32),
            [256, 256, 256, 256],
            [1024],
            norm_layer=nn.BatchNorm2d,
        )

        self.representation_size = 1024
        box_predictor = FastRCNNPredictor(self.representation_size, self.num_classes)

        self.roi_heads = RoIHeads(
            box_roi_pool,
            box_head,
            box_predictor,
            box_fg_iou_thresh,
            box_bg_iou_thresh,
            box_batch_size_per_image,
            box_positive_fraction,
            bbox_reg_weights,
            box_score_thresh,
            box_nms_thresh,
            box_detections_per_img,
            export_mode=self.export_mode,
        )
