from transformers.models.detr.modeling_detr import DetrObjectDetectionOutput
from transformers.models.detr.modeling_detr import build_position_encoding
from transformers.models.detr.modeling_detr import DetrDecoder
from transformers.models.detr.modeling_detr import DetrMLPPredictionHead

""""
transformers library changed the way detr loss was handled in 4.46.2, which was
originally working in 4.37.2
"""

# try:
#     from transformers.models.detr.modeling_detr import DetrHungarianMatcher
#     from transformers.models.detr.modeling_detr import DetrLoss
# except ImportError:
from geowatch.tasks.fusion.methods._vendor.modeling_detr import DetrHungarianMatcher
from geowatch.tasks.fusion.methods._vendor.modeling_detr import DetrLoss

from transformers import DetrConfig
import torch
import torch.nn as nn
from typing import List, Optional, Tuple, Union
from einops import rearrange


class DetrDecoderForObjectDetection(nn.Module):
    """
    Ignore:
        >>> from geowatch.tasks.fusion.methods.object_head import *  # NOQA
        >>> import torch
        >>> import kwimage
        >>> batch_size = 3
        >>> inputs = torch.rand(batch_size, 14, 14, 256)
        >>> config = DetrConfig()
        >>> self = DetrDecoderForObjectDetection(config, 256)
        >>> #
        >>> labels = []
        >>> for batch_idx in range(batch_size):
        >>>     num_item_truth = 0
        >>>     true_boxes = kwimage.Boxes.random(num_item_truth).tensor()
        >>>     labels.append({
        >>>         'class_labels': torch.LongTensor([0] * len(true_boxes)),
        >>>         'boxes': true_boxes.to_cxywh().data.float()
        >>>     })
        >>> output = self.forward(inputs, labels=labels, return_dict=True)
        >>> print(f'output.loss = {ub.urepr(output.loss, nl=1)}')
        >>> # A set of boxes for each batch item
        >>> out_boxes = kwimage.Boxes(output.pred_boxes, 'cxywh')
    """
    def __init__(self, config: DetrConfig, d_model=10, d_hidden=100):
        super().__init__()
        self.config = config
        self.object_queries = build_position_encoding(config)
        self.query_position_embeddings = nn.Embedding(config.num_queries, d_model)
        self.decoder = DetrDecoder(config)

        self.bbox_predictor = DetrMLPPredictionHead(
            input_dim=d_model, hidden_dim=100, output_dim=4, num_layers=3
        )

        self.class_labels_classifier = nn.Linear(
            config.d_model, config.num_labels + 1
        )

    # taken from https://github.com/facebookresearch/detr/blob/master/models/detr.py
    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

    def forward(
        self,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[List[dict]] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        loss_only: Optional[bool] = None,
        pred_boxes: Optional[torch.FloatTensor] = None,
        logits: Optional[torch.FloatTensor] = None
    ) -> Union[Tuple[torch.FloatTensor], DetrObjectDetectionOutput]:

        if not loss_only:
            #return_dict = return_dict if return_dict is not None else self.config.use_return_dict
            pixel_values = inputs_embeds.permute(0, 3, 1, 2)
            batch_size, num_channels, height, width = pixel_values.shape
            pixel_mask = torch.ones(((batch_size, height, width))).to(pixel_values.device)

            object_queries = self.object_queries(pixel_values, pixel_mask)

            object_queries = object_queries.flatten(2).permute(0, 2, 1)

            query_position_embeddings = self.query_position_embeddings.weight.unsqueeze(0).repeat(batch_size, 1, 1)
            queries = torch.zeros_like(query_position_embeddings)

            inputs_embeds = rearrange(inputs_embeds, 'b h w c -> b (h w) c')

            decoder_outputs = self.decoder(
                inputs_embeds=queries,
                attention_mask=None,
                object_queries=object_queries,
                query_position_embeddings=query_position_embeddings,
                encoder_hidden_states=inputs_embeds,
            )

            sequence_output = decoder_outputs[0]

            # class logits + predicted bounding boxes
            logits = self.class_labels_classifier(sequence_output)
            pred_boxes = self.bbox_predictor(sequence_output).sigmoid()

        loss, loss_dict, auxiliary_outputs = None, None, None
        if labels is not None:
            # First: create the matcher
            matcher = DetrHungarianMatcher(
                class_cost=self.config.class_cost, bbox_cost=self.config.bbox_cost, giou_cost=self.config.giou_cost
            )
            # Second: create the criterion
            losses = ["labels", "boxes", "cardinality"]
            criterion = DetrLoss(
                matcher=matcher,
                num_classes=self.config.num_labels,
                eos_coef=self.config.eos_coefficient,
                losses=losses,
            )
            criterion.to(pred_boxes.device)
            # Third: compute the losses, based on outputs and labels
            outputs_loss = {}
            outputs_loss["logits"] = logits
            outputs_loss["pred_boxes"] = pred_boxes

            if self.config.auxiliary_loss:
                # broken?
                raise NotImplementedError
                # intermediate = outputs.intermediate_hidden_states if return_dict else outputs[4]
                # outputs_class = self.class_labels_classifier(intermediate)
                # outputs_coord = self.bbox_predictor(intermediate).sigmoid()
                # auxiliary_outputs = self._set_aux_loss(outputs_class, outputs_coord)
                # outputs_loss["auxiliary_outputs"] = auxiliary_outputs

            loss_dict = criterion(outputs_loss, labels)
            # Fourth: compute total loss, as a weighted sum of the various losses
            weight_dict = {"loss_ce": 1, "loss_bbox": self.config.bbox_loss_coefficient}
            weight_dict["loss_giou"] = self.config.giou_loss_coefficient
            if self.config.auxiliary_loss:
                aux_weight_dict = {}
                for i in range(self.config.decoder_layers - 1):
                    aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
                weight_dict.update(aux_weight_dict)
            loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)

        if not return_dict:
            if auxiliary_outputs is not None:
                output = (pred_boxes) + auxiliary_outputs
            else:
                output = (pred_boxes, logits)
            return ((loss, loss_dict)) + output if loss is not None else output

        return DetrObjectDetectionOutput(
            loss=loss,
            loss_dict=loss_dict,
            logits=logits,
            pred_boxes=pred_boxes,
        )
