# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import torch

from models.backbone import Backbone, Joiner
from models.detr import DETR, PostProcess
from models.position_encoding import PositionEmbeddingSine
from models.segmentation import DETRsegm, PostProcessPanoptic
from models.transformer import Transformer

dependencies = ["torch", "torchvision"]


def _make_detr(backbone_name: str, dilation=False, num_classes=91, mask=False):
    hidden_dim = 256
    backbone = Backbone(backbone_name, train_backbone=True, return_interm_layers=mask, dilation=dilation)
    pos_enc = PositionEmbeddingSine(hidden_dim // 2, normalize=True)
    backbone_with_pos_enc = Joiner(backbone, pos_enc)
    backbone_with_pos_enc.num_channels = backbone.num_channels
    transformer = Transformer(d_model=hidden_dim, return_intermediate_dec=True)
    detr = DETR(backbone_with_pos_enc, transformer, num_classes=num_classes, num_queries=100)
    if mask:
        return DETRsegm(detr)
    return detr


def detr_resnet50(pretrained=False, num_classes=91, return_postprocessor=False):
    """
    DETR R50 with 6 encoder and 6 decoder layers.

    Achieves 42/62.4 AP/AP50 on COCO val5k.
    """
    model = _make_detr("resnet50", dilation=False, num_classes=num_classes)
    if pretrained:
        checkpoint = torch.hub.load_state_dict_from_url(
            url="https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth", map_location="cpu", check_hash=True
        )
        model.load_state_dict(checkpoint["model"])
    if return_postprocessor:
        return model, PostProcess()
    return model


def detr_resnet50_dc5(pretrained=False, num_classes=91, return_postprocessor=False):
    """
    DETR-DC5 R50 with 6 encoder and 6 decoder layers.

    The last block of ResNet-50 has dilation to increase
    output resolution.
    Achieves 43.3/63.1 AP/AP50 on COCO val5k.
    """
    model = _make_detr("resnet50", dilation=True, num_classes=num_classes)
    if pretrained:
        checkpoint = torch.hub.load_state_dict_from_url(
            url="https://dl.fbaipublicfiles.com/detr/detr-r50-dc5-f0fb7ef5.pth", map_location="cpu", check_hash=True
        )
        model.load_state_dict(checkpoint["model"])
    if return_postprocessor:
        return model, PostProcess()
    return model


def detr_resnet101(pretrained=False, num_classes=91, return_postprocessor=False):
    """
    DETR-DC5 R101 with 6 encoder and 6 decoder layers.

    Achieves 43.5/63.8 AP/AP50 on COCO val5k.
    """
    model = _make_detr("resnet101", dilation=False, num_classes=num_classes)
    if pretrained:
        checkpoint = torch.hub.load_state_dict_from_url(
            url="https://dl.fbaipublicfiles.com/detr/detr-r101-2c7b67e5.pth", map_location="cpu", check_hash=True
        )
        model.load_state_dict(checkpoint["model"])
    if return_postprocessor:
        return model, PostProcess()
    return model


def detr_resnet101_dc5(pretrained=False, num_classes=91, return_postprocessor=False):
    """
    DETR-DC5 R101 with 6 encoder and 6 decoder layers.

    The last block of ResNet-101 has dilation to increase
    output resolution.
    Achieves 44.9/64.7 AP/AP50 on COCO val5k.
    """
    model = _make_detr("resnet101", dilation=True, num_classes=num_classes)
    if pretrained:
        checkpoint = torch.hub.load_state_dict_from_url(
            url="https://dl.fbaipublicfiles.com/detr/detr-r101-dc5-a2e86def.pth", map_location="cpu", check_hash=True
        )
        model.load_state_dict(checkpoint["model"])
    if return_postprocessor:
        return model, PostProcess()
    return model


def detr_resnet50_panoptic(
    pretrained=False, num_classes=250, threshold=0.85, return_postprocessor=False
):
    """
    DETR R50 with 6 encoder and 6 decoder layers.
    Achieves 43.4 PQ on COCO val5k.

   threshold is the minimum confidence required for keeping segments in the prediction
    """
    model = _make_detr("resnet50", dilation=False, num_classes=num_classes, mask=True)
    is_thing_map = {i: i <= 90 for i in range(250)}
    if pretrained:
        checkpoint = torch.hub.load_state_dict_from_url(
            url="https://dl.fbaipublicfiles.com/detr/detr-r50-panoptic-00ce5173.pth",
            map_location="cpu",
            check_hash=True,
        )
        model.load_state_dict(checkpoint["model"])
    if return_postprocessor:
        return model, PostProcessPanoptic(is_thing_map, threshold=threshold)
    return model


def detr_resnet50_dc5_panoptic(
    pretrained=False, num_classes=250, threshold=0.85, return_postprocessor=False
):
    """
    DETR-DC5 R50 with 6 encoder and 6 decoder layers.

    The last block of ResNet-50 has dilation to increase
    output resolution.
    Achieves 44.6 on COCO val5k.

   threshold is the minimum confidence required for keeping segments in the prediction
    """
    model = _make_detr("resnet50", dilation=True, num_classes=num_classes, mask=True)
    is_thing_map = {i: i <= 90 for i in range(250)}
    if pretrained:
        checkpoint = torch.hub.load_state_dict_from_url(
            url="https://dl.fbaipublicfiles.com/detr/detr-r50-dc5-panoptic-da08f1b1.pth",
            map_location="cpu",
            check_hash=True,
        )
        model.load_state_dict(checkpoint["model"])
    if return_postprocessor:
        return model, PostProcessPanoptic(is_thing_map, threshold=threshold)
    return model


def detr_resnet101_panoptic(
    pretrained=False, num_classes=250, threshold=0.85, return_postprocessor=False
):
    """
    DETR-DC5 R101 with 6 encoder and 6 decoder layers.

    Achieves 45.1 PQ on COCO val5k.

   threshold is the minimum confidence required for keeping segments in the prediction
    """
    model = _make_detr("resnet101", dilation=False, num_classes=num_classes, mask=True)
    is_thing_map = {i: i <= 90 for i in range(250)}
    if pretrained:
        checkpoint = torch.hub.load_state_dict_from_url(
            url="https://dl.fbaipublicfiles.com/detr/detr-r101-panoptic-40021d53.pth",
            map_location="cpu",
            check_hash=True,
        )
        model.load_state_dict(checkpoint["model"])
    if return_postprocessor:
        return model, PostProcessPanoptic(is_thing_map, threshold=threshold)
    return model