import numpy as np 
from transformers.models.deformable_detr.modeling_deformable_detr import DeformableDetrMLPPredictionHead
import torch.nn as nn
import torch
def PairDetr(model, num_queries, num_classes):
    in_features = model.class_embed[0].in_features        
    model.model.query_position_embeddings = nn.Embedding(num_queries, 512)
    class_embed = nn.Linear(in_features, num_classes)
    bbox_embed = DeformableDetrMLPPredictionHead(
        input_dim=256, hidden_dim=256, output_dim=8, num_layers=3
    )
    model.class_embed = nn.ModuleList([class_embed for _ in range(6)])
    model.bbox_embed  = nn.ModuleList([bbox_embed for _ in range(6)])
    return model
    
def inverse_sigmoid(x, eps=1e-5):
    x = x.clamp(min=0, max=1)
    x1 = x.clamp(min=eps)
    x2 = (1 - x).clamp(min=eps)
    return torch.log(x1 / x2)
    
def forward(model,
        pixel_values,
        pixel_mask=None,
        decoder_attention_mask=None,
        encoder_outputs=None,
        inputs_embeds=None,
        decoder_inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,) -> torch.Tensor:
        return_dict = return_dict if return_dict is not None else model.config.use_return_dict

        outputs = model.model(
            pixel_values,
            pixel_mask=pixel_mask,
            decoder_attention_mask=decoder_attention_mask,
            encoder_outputs=encoder_outputs,
            inputs_embeds=inputs_embeds,
            decoder_inputs_embeds=decoder_inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = outputs.intermediate_hidden_states if return_dict else outputs[2]
        init_reference = outputs.init_reference_points if return_dict else outputs[0]
        inter_references = outputs.intermediate_reference_points if return_dict else outputs[3]
        outputs_classes = []
        outputs_coords = []
        cons = inverse_sigmoid(init_reference)
        for level in range(hidden_states.shape[1]):
            if level == 0:
                reference = init_reference
            else:
                reference = inter_references[:, level - 1]
            reference = inverse_sigmoid(reference)
            outputs_class = model.class_embed[level](hidden_states[:, level])
            delta_bbox = model.bbox_embed[level](hidden_states[:, level])
            if reference.shape[-1] == 4:
                delta_bbox[..., :4] += reference
                outputs_coord_logits = delta_bbox
            elif reference.shape[-1] == 2:
                delta_bbox[..., :2] += reference
                delta_bbox[..., 4:6] += cons
                outputs_coord_logits = delta_bbox
            else:
                raise ValueError(f"reference.shape[-1] should be 4 or 2, but got {reference.shape[-1]}")
            outputs_coord = outputs_coord_logits.sigmoid()
            outputs_classes.append(outputs_class)
            outputs_coords.append(outputs_coord)
        outputs_class = torch.stack(outputs_classes, dim=1)
        outputs_coord = torch.stack(outputs_coords, dim=1)

        logits = outputs_class[:, -1]
        pred_boxes = outputs_coord[:, -1]

        dict_outputs = {
            "logits":logits,
            "pred_boxes": pred_boxes,
            "init_reference_points": outputs.init_reference_points,
                       }    
        return dict_outputs