# ------------------------------------------------------------------------
# Deformable DETR
# Copyright (c) 2020 SenseTime. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Modified from DETR (https://github.com/facebookresearch/detr)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# ------------------------------------------------------------------------

"""
Deformable DETR model and criterion classes.
"""
import torch
import torch.nn.functional as F
from torch import nn
import math

from util import box_ops
from util.misc import (NestedTensor, nested_tensor_from_tensor_list,
                       accuracy, get_world_size, interpolate,
                       is_dist_avail_and_initialized, inverse_sigmoid)

from .backbone import build_backbone, build_backbone_tea
from .matcher import build_matcher
from .segmentation import (DETRsegm, PostProcessPanoptic, PostProcessSegm,
                           dice_loss, sigmoid_focal_loss)
from .deformable_transformer import build_deforamble_transformer
import copy


def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

import torch
import torch.nn as nn
import torch.nn.functional as F


class MyScaleDotProductAttention(nn.Module):
    """
    compute scale dot product attention

    Query : given sentence that we focused on (decoder)
    Key : every sentence to check relationship with Qeury(encoder)
    Value : every sentence same with Key (encoder)
    """

    def __init__(self):
        super(MyScaleDotProductAttention, self).__init__()
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(0.0)

    def forward(self, q, k, v, mask=None, e=1e-12, tempeature=1.0):
        # input is 4 dimension tensor
        # [batch_size, head, length, d_tensor]
        batch_size, head, length, d_tensor = k.size()

        # 1. dot product Query with Key^T to compute similarity
        k_t = k.transpose(2, 3)  # transpose
        score = (q @ k_t) / math.sqrt(d_tensor)  # scaled dot product

        # 2. apply masking (opt)
        if mask is not None:
            score = score.masked_fill(mask == 0, -10000)

        # 3. pass them softmax to make [0, 1] range
        score = self.softmax(score / tempeature)

        score = self.dropout(score)
        # 4. multiply with Value
        v = score @ v

        if score.shape[1] != 1:
            print(score.shape, k.size(), " not equal 1 !!!!")

        return v, score.sum(dim=1)


class MyMultiHeadAttention(nn.Module):

    def __init__(self, d_model, n_head):
        super(MyMultiHeadAttention, self).__init__()
        self.n_head = n_head
        self.attention = MyScaleDotProductAttention()
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_concat = nn.Linear(d_model, d_model)

    def forward(self, query, key, value, mask=None, tempeature=1.0):
        q = query
        k = key
        v = value
        # 1. dot product with weight matrices
        q, k, v = self.w_q(q), self.w_k(k), self.w_v(v)

        # 2. split tensor by number of heads
        q, k, v = self.split(q), self.split(k), self.split(v)

        # 3. do scale dot product to compute similarity
        out, attention = self.attention(q, k, v, mask=mask, tempeature=tempeature)

        # 4. concat and pass to linear layer
        out = self.concat(out)
        out = self.w_concat(out)

        # 5. visualize attention map
        # TODO : we should implement visualization

        return out, attention

    def split(self, tensor):
        """
        split tensor by number of head

        :param tensor: [batch_size, length, d_model]
        :return: [batch_size, head, length, d_tensor]
        """
        batch_size, length, d_model = tensor.size()

        d_tensor = d_model // self.n_head
        tensor = tensor.view(batch_size, length, self.n_head, d_tensor).transpose(1, 2)
        # it is similar with group convolution (split by number of heads)

        return tensor

    def concat(self, tensor):
        """
        inverse function of self.split(tensor : torch.Tensor)

        :param tensor: [batch_size, head, length, d_tensor]
        :return: [batch_size, length, d_model]
        """
        batch_size, head, length, d_tensor = tensor.size()
        d_model = head * d_tensor

        tensor = tensor.transpose(1, 2).contiguous().view(batch_size, length, d_model)
        return tensor


class MyCrossAttentionBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, ffn_dim, dropout=0.0):
        super().__init__()
        #self.cross_attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, bias=False, batch_first=True)
        self.cross_attn = MyMultiHeadAttention(embed_dim, num_heads)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.dropout1 = nn.Dropout(dropout)

        self.ffn = nn.Sequential(
            nn.Linear(embed_dim, ffn_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, embed_dim)
        )
        self.norm2 = nn.LayerNorm(embed_dim)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, tgt, memory, value,  memory_mask=None, tgt_mask=None, return_early=False, tempeature=1.0):
        """
        tgt:    [B, T_tgt, E] -- query input (e.g., decoder state)
        memory: [B, T_mem, E] -- context input (e.g., encoder output)
        memory_mask: [B, T_tgt, T_mem]
        """
        # Cross-attention: tgt attends to memory
        attn_output, ws = self.cross_attn(
            query=tgt,
            key=memory,
            value=value,
            tempeature=tempeature,
        )

        if return_early:
            return tgt, attn_output, ws

        tgt = tgt + self.dropout1(attn_output)
        tgt = self.norm1(tgt)
        # res = tgt + self.dropout1(attn_output)
        # res = self.norm1(res)

        # Feed-forward network
        ffn_output = self.ffn(tgt)
        tgt = tgt + self.dropout2(ffn_output)
        tgt = self.norm2(tgt)

        return tgt, attn_output, ws


class CrossAttentionBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, ffn_dim, dropout=0.0):
        super().__init__()
        self.cross_attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, bias=False, batch_first=True)
        #self.cross_attn = nn.MultiHeadAttention(embed_dim, num_heads)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.dropout1 = nn.Dropout(dropout)

        self.ffn = nn.Sequential(
            nn.Linear(embed_dim, ffn_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, embed_dim)
        )
        self.norm2 = nn.LayerNorm(embed_dim)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, tgt, memory, value,  memory_mask=None, tgt_mask=None, return_early=False, tempeature=1.0):
        """
        tgt:    [B, T_tgt, E] -- query input (e.g., decoder state)
        memory: [B, T_mem, E] -- context input (e.g., encoder output)
        memory_mask: [B, T_tgt, T_mem]
        """
        # Cross-attention: tgt attends to memory
        attn_output, ws = self.cross_attn(
            query=tgt,
            key=memory,
            value=value,
            key_padding_mask=memory_mask,  # optional mask
            attn_mask=tgt_mask             # optional temporal mask
        )

        if return_early:
            return tgt, attn_output, ws

        tgt = tgt + self.dropout1(attn_output)
        tgt = self.norm1(tgt)
        # res = tgt + self.dropout1(attn_output)
        # res = self.norm1(res)

        # Feed-forward network
        ffn_output = self.ffn(tgt)
        tgt = tgt + self.dropout2(ffn_output)
        tgt = self.norm2(tgt)

        return tgt, attn_output, ws


class Auto_match(nn.Module):
    def __init__(self, para=None):
        super().__init__()
        
        mults = 1
        self.mults = mults
        self.cc = 256
        self.heads = 1

        self.embed = nn.Embedding(256, self.cc*mults)

        self.enc1 = MyCrossAttentionBlock(self.cc*mults, self.heads, self.cc*mults)
        self.enc2 = MyCrossAttentionBlock(self.cc*mults, self.heads, self.cc*mults)
        self.enc3 = MyCrossAttentionBlock(self.cc*mults, self.heads, self.cc*mults)

        self.self_enc1 = MyCrossAttentionBlock(self.cc*mults, self.heads, self.cc*mults)
        self.self_enc2 = MyCrossAttentionBlock(self.cc*mults, self.heads, self.cc*mults)

        self.dec1 = MyCrossAttentionBlock(self.cc*mults, self.heads, self.cc*mults)
        self.dec2 = MyCrossAttentionBlock(self.cc*mults, self.heads, self.cc*mults)
        self.dec3 = MyCrossAttentionBlock(self.cc*mults, self.heads, self.cc*mults)

        self.self_dec1 = MyCrossAttentionBlock(self.cc*mults, self.heads, self.cc*mults)
        self.self_dec2 = MyCrossAttentionBlock(self.cc*mults, self.heads, self.cc*mults)

        self.ffn_enc_query = nn.Sequential(
                nn.Linear(95, mults*self.cc),
                nn.ReLU(inplace=True),
                nn.Linear(mults*self.cc, mults*self.cc),
                nn.LayerNorm(self.cc*mults),
            )

        self.ffn_enc_key = nn.Sequential(
                nn.Linear(95, mults*self.cc),
                nn.ReLU(inplace=True),
                nn.Linear(mults*self.cc, mults*self.cc),
                nn.LayerNorm(self.cc*mults),
            )

        self.ffn_enc_value = nn.Sequential(
                nn.Linear(95, mults*self.cc),
                nn.ReLU(inplace=True),
                nn.Linear(mults*self.cc, mults*self.cc),
                nn.LayerNorm(self.cc*mults),
            )


        self.ffn_fusion_enc = nn.Sequential(
                nn.Linear(self.cc*mults * 1, self.cc*mults * 2),
                nn.ReLU(inplace=True),
                nn.Linear(self.cc*mults * 2, self.cc*mults),
                nn.LayerNorm(self.cc*mults),
            )

        self.ffn_fusion_dec = nn.Sequential(
                nn.Linear(self.cc*mults * 1, self.cc*mults * 2),
                nn.LayerNorm(self.cc*mults * 2),
                nn.ReLU(inplace=True),
                nn.Linear(self.cc*mults * 2, self.cc*mults * 2),
                nn.ReLU(inplace=True),
                nn.Linear(self.cc*mults * 2, self.cc),
            )

        self.lyn_emb = nn.LayerNorm(mults*self.cc)

        self.counter = 0

    def enc_dec(self, hs, hs_aug, embed):
        hs_adp_query = self.ffn_enc_query(hs)
        hs_adp_key = self.ffn_enc_key(hs_aug)
        hs_adp_value = self.ffn_enc_value(hs_aug)
        out_gt, o1, ww = self.enc1(hs_adp_query, hs_adp_key, hs_adp_value, return_early=True)

        o1 = torch.bmm(ww, hs_aug)
        
        return o1, hs_adp_query, hs_adp_key, hs_adp_value, ww

    def only_dec(self, hs_adp_tea, hs_adp_stu, hs_adp_value, hs_stu):
        out_gt, o1, ww = self.enc1(hs_adp_tea, hs_adp_stu, hs_adp_value, hs_stu, return_early=True, tempeature=4.)
        o1 = torch.bmm(ww, hs_stu)
        return o1, ww

    def norm(self, x):
        #x = (x - torch.mean(x, dim=-1, keepdim=True)) / (torch.std(x, dim=-1, keepdim=True) + 1e-5)
        return x


    def high_inter_align(self, fea_tea, fea_stu, module):
        # feature: bs, n, c
        feature = fea_tea #torch.cat([fea_tea, fea_stu], dim=1)
        
        fea_sp = feature.shape
        weight = torch.rand(fea_sp[0], fea_sp[1], device=fea_tea.device)
        feature_bbox = feature[:,:,:4]
        feature_cls = feature[:,:,4:]

        fea_bbox_inter = feature_bbox.unsqueeze(-2) * weight.unsqueeze(-1).unsqueeze(-1) + feature_bbox.unsqueeze(-3) * (1. - weight.unsqueeze(-2).unsqueeze(-1))

        feature_inter = torch.cat([fea_bbox_inter, feature_cls.unsqueeze(-2).repeat(1,1,fea_sp[1],1)], dim=-1)
        
        feature_aug = torch.cat([feature_bbox.unsqueeze(-3).repeat(1,fea_sp[1],1,1), feature_cls.unsqueeze(-2).repeat(1,1,fea_sp[1],1)], dim=-1)
        high_fea = module(feature_aug)
        feature_org = module(feature)
        high_fea_inter_tmp = feature_org.unsqueeze(-2) * weight.unsqueeze(-1).unsqueeze(-1) + high_fea * (1. - weight.unsqueeze(-2).unsqueeze(-1))
        high_fea_inter = module(feature_inter)

        #print(fea_bbox_inter.shape, feature_inter.shape, feature_aug.shape, high_fea.shape, high_fea_inter_tmp.shape)

        loss = torch.mean(torch.abs(high_fea_inter - high_fea_inter_tmp))

        return loss

    def forward(self, hs_stu, hs_tea, costs_stu, costs_tea, opt_nom=True):
        ssp = hs_stu.shape
        embed = self.lyn_emb(self.embed.weight.unsqueeze(0).repeat(ssp[0],1,1))
        #embed_detail = self.lyn_emb_detail(self.embed_detail.weight.unsqueeze(0).repeat(ssp[0],1,1))

        if costs_stu is not None and costs_tea is not None:
            if opt_nom:
                self.enc1.cross_attn.attention.dropout.p = 0.0
            else:
                self.enc1.cross_attn.attention.dropout.p = 0.

            #tmp_mask = torch.ones_like(hs_tea, device=hs_tea.device)
            #tmp_mask[:,:,4:] = tmp_mask[:,:,4:] * 0.

            topks = 50
            w_c_tea = torch.exp(-torch.clamp(costs_tea*.5, min=-1., max=5.))    
            w_c_stu = torch.exp(-torch.clamp(costs_stu*.5, min=-1., max=5.))
            out_tmp = torch.topk(-costs_tea, k=topks, dim=1)
            out_stu_tmp = torch.topk(-costs_stu, k=topks, dim=1)
            hs_tea = hs_tea[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1)]
            hs_stu = hs_stu[torch.arange(ssp[0]).view(-1, 1), out_stu_tmp[1].squeeze(-1)]
            w_c_tea = w_c_tea[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1)]
            w_c_stu = w_c_stu[torch.arange(ssp[0]).view(-1, 1), out_stu_tmp[1].squeeze(-1)]

            re_tea, hs_adp_tea_query, hs_adp_tea_key, hs_adp_tea_value, ws1 = self.enc_dec(hs_tea, torch.cat([hs_tea, hs_stu], dim=1), embed.clone())
            re_stu, hs_adp_stu_query, hs_adp_stu_key, hs_adp_stu_value, ws2 = self.enc_dec(hs_stu, hs_stu, embed.clone())
            
            #hs_adp_tea_query = hs_adp_tea_query[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1)]

            if not opt_nom:
                hs_adp_tea_key = hs_adp_tea_key[:,:topks,:]#[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1)]
                hs_adp_tea_value = hs_adp_tea_value[:,:topks,:]#[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1)]
                #hs_tea = hs_tea[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1)]

                #w_c_tea = w_c_tea[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1)]

                #out_stu_tmp = torch.topk(-costs_stu, k=50, dim=1)
                hs_adp_stu_key = hs_adp_stu_key#[torch.arange(ssp[0]).view(-1, 1), out_stu_tmp[1].squeeze(-1)]
                hs_adp_stu_value = hs_adp_stu_value#[torch.arange(ssp[0]).view(-1, 1), out_stu_tmp[1].squeeze(-1)]
                #hs_stu = hs_stu[torch.arange(ssp[0]).view(-1, 1), out_stu_tmp[1].squeeze(-1)]
                
            if not opt_nom:
                re_stu_oo1, www1 = self.only_dec(hs_adp_stu_query.detach(), hs_adp_tea_key, hs_adp_tea_value, hs_tea)
            else:
                re_stu_oo1, www1 = self.only_dec(hs_adp_stu_query.detach(), hs_adp_tea_key[:,:topks,:], hs_adp_tea_value[:,:topks,:], hs_tea)
            re_stu_oo2, www2 = self.only_dec(hs_adp_stu_query.detach(), hs_adp_stu_key, hs_adp_stu_value, hs_stu)
            #re_stu_oo2, www2 = self.only_dec(hs_adp_tea_query.detach(), hs_adp_tea_key, hs_adp_tea_value, hs_stu)

            #ws1_tmp = ws1[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1)]
            #ws1 = ws1[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1), out_tmp[1].squeeze(-1)]

            #re_tea = re_tea[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1)]
            #if opt_nom:
            #    hs_tea = hs_tea[torch.arange(ssp[0]).view(-1, 1), out_tmp[1].squeeze(-1)]

            # inter_loss = self.high_inter_align(hs_tea, hs_tea, self.ffn_enc_query) + \
            #              self.high_inter_align(hs_tea, hs_tea, self.ffn_enc_key)

            reg_cls_ww = torch.ones_like(hs_tea, device=hs_tea.device)
            reg_cls_ww[:,:,:4] = reg_cls_ww[:,:,:4] * 2.
            reg_cls_ww2 = torch.ones_like(hs_tea, device=hs_tea.device)
            reg_cls_ww2[:,:,:4] = reg_cls_ww2[:,:,:4] * 0. + 1.

            #print(www1.shape, w_c_tea.shape, " new !!!" * 3)
            wwww = torch.ones_like(torch.bmm(www1, w_c_tea) + 1e-1, device=w_c_tea.device) #torch.ones_like(w_c_tea, device=w_c_tea.device) #torch.bmm(www1, w_c_tea)
            # print(wwww.shape, re_stu_oo1.shape, " < > >DD " * 5)
            # dis_loss = torch.mean(torch.abs(info_stu - info_tea.detach()))
            d1 = torch.abs(re_stu_oo1 - re_stu_oo2)
            # cos_loss = torch.mean((1.0 - torch.clamp(F.cosine_similarity(self.norm(re_stu_oo1[:,:,4:]), self.norm(re_stu_oo2[:,:,4:]), dim=-1), min=-1., max=1.)) * wwww.squeeze(-1)) / torch.mean(wwww)
            dis_loss = torch.mean(d1 * reg_cls_ww2 * reg_cls_ww * wwww) / torch.mean(wwww)
                       
            #dis_loss = torch.mean((1.0 - F.cosine_similarity(re_stu_o, re_tea.detach(), dim=-1)) * w_c_tea.squeeze(-1))
            # ww_tea = torch.pow(torch.abs(re_tea - hs_tea.detach()) + 1e-8, 0.5)
            # ww_stu = torch.pow(torch.abs(re_stu - hs_stu.detach()) + 1e-8, 0.5)
            recons_loss = torch.mean(torch.abs(re_tea - hs_tea.detach()) * reg_cls_ww* reg_cls_ww) #+ \
            #            torch.mean(torch.abs(re_stu - hs_stu.detach()) * reg_cls_ww* reg_cls_ww)
            
            #ww_loss = torch.mean(torch.abs(torch.sum(www1, dim=1) - torch.sum(www2, dim=1)))
            labs = torch.eye(topks, device=hs_tea.device).unsqueeze(0) 
            ws1 = ws1[:,:,:topks]
            if opt_nom:
                entropy_loss = - torch.mean(torch.sum(labs * torch.log(torch.clamp(ws1, min=-1e-8, max=0.9)+1e-8), dim=-1) ) 
                            #    + \
                            #    + torch.mean(torch.sum(ws1_tmp*torch.log(ws1_tmp+1e-8), dim=-1)) * 0.2
            else:
                # entropy_loss = -torch.mean( torch.sum(www1 * torch.log(www1 + 1e-8), dim=-1) ) * 1.
                entropy_loss = - torch.mean(torch.sum(labs * torch.log(torch.clamp(ws1, min=-1e-8, max=0.9)+1e-8), dim=-1) ) 
                            #    + \
                            #    + torch.mean(torch.sum(ws1_tmp*torch.log(ws1_tmp+1e-8), dim=-1)) * 0.2
            self.counter += 1
            if self.counter % 100 < 4:
                print(torch.mean(torch.sum(ws1 * torch.log(ws1+1e-8), dim=-1)), torch.mean(torch.sum(www1 * torch.log(www1+1e-8), dim=-1)), " < > " * 3)

            
            losses = {}
            losses['nom_dis_loss'] = dis_loss * 20.
            losses['recons_loss'] = recons_loss * 1.
            losses['ww_dist_loss'] = recons_loss * 0.
            losses['entropy_loss'] = entropy_loss * 2.
            #losses['high_inter_loss'] = inter_loss * 2.

        else:
            re_tea, info_tea_query, info_tea_key, _, _ = self.enc_dec(hs_tea, hs_tea, embed.clone())
            re_stu, info_stu_query, info_stu_key, _, _ = self.enc_dec(hs_stu, hs_stu, embed.clone())
            #re_stu, info_stu = self.enc_dec(hs_stu, embed.clone())


            dis_loss = torch.mean(torch.abs(info_stu_query - info_tea_query.detach()))
            ww_tea = torch.pow(torch.abs(re_tea - hs_tea.detach()) + 1e-8, 0.5)
            ww_stu = torch.pow(torch.abs(re_stu - hs_stu.detach()) + 1e-8, 0.5)
            recons_loss = torch.mean(torch.abs(re_tea - hs_tea.detach()) * ww_tea) + \
                        torch.mean(torch.abs(re_stu - hs_stu.detach()) * ww_stu)
            
            # recons_loss = torch.mean(1.0 - F.cosine_similarity(re_stu, hs_stu.detach(), dim=-1)) + \
            #               torch.mean(1.0 - F.cosine_similarity(re_tea, hs_tea.detach(), dim=-1))

            losses = {}
            losses['nom_dis_loss'] = dis_loss * 0.
            losses['recons_loss'] = recons_loss * 0.
            losses['ww_dist_loss'] = recons_loss * 0.
            losses['entropy_loss'] = recons_loss * 0.
            losses['high_inter_loss'] = recons_loss * 0.

        return re_stu, losses
  

class Auto_match_layers(nn.Module):
    def __init__(self, layers_num=1):
        super().__init__()

        self.layers = nn.ModuleList()
        for i in range(layers_num):
            self.layers.append(Auto_match())
        
    def forward(self, hs, hs_tea, costs_stu, costs, opt_nom=True):
        re_stu1, losses1 = self.layers[0](hs[0,...], hs_tea[2,...].detach(), costs_stu[1], costs[3], opt_nom)
        re_stu2, losses2 = self.layers[1](hs[1,...], hs_tea[2,...].detach(), costs_stu[2], costs[3], opt_nom)

        if opt_nom:
            losses1['nom_dis_loss'] = (losses1['nom_dis_loss'] + losses2['nom_dis_loss']).detach() * .0
            losses1['recons_loss'] = (losses1['recons_loss'] + losses2['recons_loss']) * 0.5
            losses1['ww_dist_loss'] = (losses1['ww_dist_loss'] + losses2['ww_dist_loss']).detach() * .0
            losses1['entropy_loss'] = (losses1['entropy_loss'] + losses2['entropy_loss']) * 1.0
        else:
            losses1['nom_dis_loss'] = (losses1['nom_dis_loss'] + losses2['nom_dis_loss']) * 1.
            losses1['recons_loss'] = (losses1['recons_loss'] + losses2['recons_loss']).detach() * 1.0
            losses1['ww_dist_loss'] = (losses1['ww_dist_loss'] + losses2['ww_dist_loss']) * 1.
            losses1['entropy_loss'] = (losses1['entropy_loss'] + losses2['entropy_loss']).detach() * 1.0
        
        return losses1

class DeformableDETR(nn.Module):
    """ This is the Deformable DETR module that performs object detection """
    def __init__(self, backbone, transformer, num_classes, num_queries, num_feature_levels,
                 aux_loss=True, with_box_refine=False, two_stage=False):
        """ Initializes the model.
        Parameters:
            backbone: torch module of the backbone to be used. See backbone.py
            transformer: torch module of the transformer architecture. See transformer.py
            num_classes: number of object classes
            num_queries: number of object queries, ie detection slot. This is the maximal number of objects
                         DETR can detect in a single image. For COCO, we recommend 100 queries.
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
            with_box_refine: iterative bounding box refinement
            two_stage: two-stage Deformable DETR
        """
        super().__init__()
        self.num_queries = num_queries
        self.transformer = transformer
        hidden_dim = transformer.d_model
        self.class_embed = nn.Linear(hidden_dim, num_classes)
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        self.num_feature_levels = num_feature_levels
        if not two_stage:
            self.query_embed = nn.Embedding(num_queries, hidden_dim*2)
        if num_feature_levels > 1:
            num_backbone_outs = len(backbone.strides)
            input_proj_list = []
            for _ in range(num_backbone_outs):
                in_channels = backbone.num_channels[_]
                input_proj_list.append(nn.Sequential(
                    nn.Conv2d(in_channels, hidden_dim, kernel_size=1),
                    nn.GroupNorm(32, hidden_dim),
                ))
            for _ in range(num_feature_levels - num_backbone_outs):
                input_proj_list.append(nn.Sequential(
                    nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=2, padding=1),
                    nn.GroupNorm(32, hidden_dim),
                ))
                in_channels = hidden_dim
            self.input_proj = nn.ModuleList(input_proj_list)
        else:
            self.input_proj = nn.ModuleList([
                nn.Sequential(
                    nn.Conv2d(backbone.num_channels[0], hidden_dim, kernel_size=1),
                    nn.GroupNorm(32, hidden_dim),
                )])
        self.backbone = backbone
        self.aux_loss = aux_loss
        self.with_box_refine = with_box_refine
        self.two_stage = two_stage

        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        self.class_embed.bias.data = torch.ones(num_classes) * bias_value
        nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
        nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)
        for proj in self.input_proj:
            nn.init.xavier_uniform_(proj[0].weight, gain=1)
            nn.init.constant_(proj[0].bias, 0)

        # if two-stage, the last class_embed and bbox_embed is for region proposal generation
        num_pred = (transformer.decoder.num_layers + 1) if two_stage else transformer.decoder.num_layers
        if with_box_refine:
            self.class_embed = _get_clones(self.class_embed, num_pred)
            self.bbox_embed = _get_clones(self.bbox_embed, num_pred)
            nn.init.constant_(self.bbox_embed[0].layers[-1].bias.data[2:], -2.0)
            # hack implementation for iterative bounding box refinement
            self.transformer.decoder.bbox_embed = self.bbox_embed
        else:
            nn.init.constant_(self.bbox_embed.layers[-1].bias.data[2:], -2.0)
            self.class_embed = nn.ModuleList([self.class_embed for _ in range(num_pred)])
            self.bbox_embed = nn.ModuleList([self.bbox_embed for _ in range(num_pred)])
            self.transformer.decoder.bbox_embed = None
        if two_stage:
            # hack implementation for two-stage
            self.transformer.decoder.class_embed = self.class_embed
            for box_embed in self.bbox_embed:
                nn.init.constant_(box_embed.layers[-1].bias.data[2:], 0.0)

    def forward(self, samples: NestedTensor):
        """ The forward expects a NestedTensor, which consists of:
               - samples.tensor: batched images, of shape [batch_size x 3 x H x W]
               - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels

            It returns a dict with the following elements:
               - "pred_logits": the classification logits (including no-object) for all queries.
                                Shape= [batch_size x num_queries x (num_classes + 1)]
               - "pred_boxes": The normalized boxes coordinates for all queries, represented as
                               (center_x, center_y, height, width). These values are normalized in [0, 1],
                               relative to the size of each individual image (disregarding possible padding).
                               See PostProcess for information on how to retrieve the unnormalized bounding box.
               - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of
                                dictionnaries containing the two above keys for each decoder layer.
        """
        if not isinstance(samples, NestedTensor):
            samples = nested_tensor_from_tensor_list(samples)
        features, pos = self.backbone(samples)

        srcs = []
        masks = []
        for l, feat in enumerate(features):
            src, mask = feat.decompose()
            srcs.append(self.input_proj[l](src))
            masks.append(mask)
            assert mask is not None
        if self.num_feature_levels > len(srcs):
            _len_srcs = len(srcs)
            for l in range(_len_srcs, self.num_feature_levels):
                if l == _len_srcs:
                    src = self.input_proj[l](features[-1].tensors)
                else:
                    src = self.input_proj[l](srcs[-1])
                m = samples.mask
                mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(torch.bool)[0]
                pos_l = self.backbone[1](NestedTensor(src, mask)).to(src.dtype)
                srcs.append(src)
                masks.append(mask)
                pos.append(pos_l)

        query_embeds = None
        if not self.two_stage:
            query_embeds = self.query_embed.weight
        hs, init_reference, inter_references, enc_outputs_class, enc_outputs_coord_unact = self.transformer(srcs, masks, pos, query_embeds)

        outputs_classes = []
        outputs_coords = []
        for lvl in range(hs.shape[0]):
            if lvl == 0:
                reference = init_reference
            else:
                reference = inter_references[lvl - 1]
            reference = inverse_sigmoid(reference)
            outputs_class = self.class_embed[lvl](hs[lvl])
            tmp = self.bbox_embed[lvl](hs[lvl])
            if reference.shape[-1] == 4:
                tmp += reference
            else:
                assert reference.shape[-1] == 2
                tmp[..., :2] += reference
            outputs_coord = tmp.sigmoid()
            outputs_classes.append(outputs_class)
            outputs_coords.append(outputs_coord)
        outputs_class = torch.stack(outputs_classes)
        outputs_coord = torch.stack(outputs_coords)

        out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1], 'pred_logits_all': outputs_class, 'pred_boxes_all': outputs_coord}
        if self.aux_loss:
            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)

        if self.two_stage:
            enc_outputs_coord = enc_outputs_coord_unact.sigmoid()
            out['enc_outputs'] = {'pred_logits': enc_outputs_class, 'pred_boxes': enc_outputs_coord}
        return out

    @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 [{'pred_logits': a, 'pred_boxes': b}
                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]


class SetCriterion(nn.Module):
    """ This class computes the loss for DETR.
    The process happens in two steps:
        1) we compute hungarian assignment between ground truth boxes and the outputs of the model
        2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
    """
    def __init__(self, num_classes, matcher, weight_dict, losses, focal_alpha=0.25):
        """ Create the criterion.
        Parameters:
            num_classes: number of object categories, omitting the special no-object category
            matcher: module able to compute a matching between targets and proposals
            weight_dict: dict containing as key the names of the losses and as values their relative weight.
            losses: list of all the losses to be applied. See get_loss for list of available losses.
            focal_alpha: alpha in Focal Loss
        """
        super().__init__()
        self.num_classes = num_classes
        self.matcher = matcher
        self.weight_dict = weight_dict
        self.losses = losses
        self.focal_alpha = focal_alpha

    def loss_labels(self, outputs, targets, indices, num_boxes, log=True):
        """Classification loss (NLL)
        targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes]
        """
        assert 'pred_logits' in outputs
        src_logits = outputs['pred_logits']

        idx = self._get_src_permutation_idx(indices)
        target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)])
        target_classes = torch.full(src_logits.shape[:2], self.num_classes,
                                    dtype=torch.int64, device=src_logits.device)
        target_classes[idx] = target_classes_o

        target_classes_onehot = torch.zeros([src_logits.shape[0], src_logits.shape[1], src_logits.shape[2] + 1],
                                            dtype=src_logits.dtype, layout=src_logits.layout, device=src_logits.device)
        target_classes_onehot.scatter_(2, target_classes.unsqueeze(-1), 1)

        target_classes_onehot = target_classes_onehot[:,:,:-1]
        loss_ce = sigmoid_focal_loss(src_logits, target_classes_onehot, num_boxes, alpha=self.focal_alpha, gamma=2) * src_logits.shape[1]
        losses = {'loss_ce': loss_ce}

        if log:
            # TODO this should probably be a separate loss, not hacked in this one here
            losses['class_error'] = 100 - accuracy(src_logits[idx], target_classes_o)[0]
        return losses

    @torch.no_grad()
    def loss_cardinality(self, outputs, targets, indices, num_boxes):
        """ Compute the cardinality error, ie the absolute error in the number of predicted non-empty boxes
        This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients
        """
        pred_logits = outputs['pred_logits']
        device = pred_logits.device
        tgt_lengths = torch.as_tensor([len(v["labels"]) for v in targets], device=device)
        # Count the number of predictions that are NOT "no-object" (which is the last class)
        card_pred = (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1).sum(1)
        card_err = F.l1_loss(card_pred.float(), tgt_lengths.float())
        losses = {'cardinality_error': card_err}
        return losses

    def loss_boxes(self, outputs, targets, indices, num_boxes):
        """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss
           targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]
           The target boxes are expected in format (center_x, center_y, h, w), normalized by the image size.
        """
        assert 'pred_boxes' in outputs
        idx = self._get_src_permutation_idx(indices)
        src_boxes = outputs['pred_boxes'][idx]
        target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0)

        loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')

        losses = {}
        losses['loss_bbox'] = loss_bbox.sum() / num_boxes

        loss_giou = 1 - torch.diag(box_ops.generalized_box_iou(
            box_ops.box_cxcywh_to_xyxy(src_boxes),
            box_ops.box_cxcywh_to_xyxy(target_boxes)))
        losses['loss_giou'] = loss_giou.sum() / num_boxes
        return losses

    def loss_masks(self, outputs, targets, indices, num_boxes):
        """Compute the losses related to the masks: the focal loss and the dice loss.
           targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]
        """
        assert "pred_masks" in outputs

        src_idx = self._get_src_permutation_idx(indices)
        tgt_idx = self._get_tgt_permutation_idx(indices)

        src_masks = outputs["pred_masks"]

        # TODO use valid to mask invalid areas due to padding in loss
        target_masks, valid = nested_tensor_from_tensor_list([t["masks"] for t in targets]).decompose()
        target_masks = target_masks.to(src_masks)

        src_masks = src_masks[src_idx]
        # upsample predictions to the target size
        src_masks = interpolate(src_masks[:, None], size=target_masks.shape[-2:],
                                mode="bilinear", align_corners=False)
        src_masks = src_masks[:, 0].flatten(1)

        target_masks = target_masks[tgt_idx].flatten(1)

        losses = {
            "loss_mask": sigmoid_focal_loss(src_masks, target_masks, num_boxes),
            "loss_dice": dice_loss(src_masks, target_masks, num_boxes),
        }
        return losses

    def _get_src_permutation_idx(self, indices):
        # permute predictions following indices
        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
        src_idx = torch.cat([src for (src, _) in indices])
        return batch_idx, src_idx

    def _get_tgt_permutation_idx(self, indices):
        # permute targets following indices
        batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
        tgt_idx = torch.cat([tgt for (_, tgt) in indices])
        return batch_idx, tgt_idx

    def get_loss(self, loss, outputs, targets, indices, num_boxes, **kwargs):
        loss_map = {
            'labels': self.loss_labels,
            'cardinality': self.loss_cardinality,
            'boxes': self.loss_boxes,
            'masks': self.loss_masks
        }
        assert loss in loss_map, f'do you really want to compute {loss} loss?'
        return loss_map[loss](outputs, targets, indices, num_boxes, **kwargs)

    def forward(self, outputs, targets):
        """ This performs the loss computation.
        Parameters:
             outputs: dict of tensors, see the output specification of the model for the format
             targets: list of dicts, such that len(targets) == batch_size.
                      The expected keys in each dict depends on the losses applied, see each loss' doc
        """
        outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs' and k != 'enc_outputs'}

        costes = []
        # Retrieve the matching between the outputs of the last layer and the targets
        # indices, C = self.matcher(outputs_without_aux, targets, time_test=True)
        indices, C = self.matcher(outputs_without_aux, targets)

        costes.append(C)
        # Compute the average number of target boxes accross all nodes, for normalization purposes
        num_boxes = sum(len(t["labels"]) for t in targets)
        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
        if is_dist_avail_and_initialized():
            torch.distributed.all_reduce(num_boxes)
        num_boxes = torch.clamp(num_boxes / get_world_size(), min=1).item()

        # Compute all the requested losses
        losses = {}
        for loss in self.losses:
            kwargs = {}
            losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes, **kwargs))

        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
        if 'aux_outputs' in outputs:
            for i, aux_outputs in enumerate(outputs['aux_outputs']):
                if i < -1: 
                    indices, C = self.matcher(aux_outputs, targets, time_test=True)
                indices, C = self.matcher(aux_outputs, targets)
                costes.append(C)
                for loss in self.losses:
                    if loss == 'masks':
                        # Intermediate masks losses are too costly to compute, we ignore them.
                        continue
                    kwargs = {}
                    if loss == 'labels':
                        # Logging is enabled only for the last layer
                        kwargs['log'] = False
                    l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_boxes, **kwargs)
                    l_dict = {k + f'_{i}': v for k, v in l_dict.items()}
                    losses.update(l_dict)

        if 'enc_outputs' in outputs:
            enc_outputs = outputs['enc_outputs']
            bin_targets = copy.deepcopy(targets)
            for bt in bin_targets:
                bt['labels'] = torch.zeros_like(bt['labels'])
            indices = self.matcher(enc_outputs, bin_targets)
            for loss in self.losses:
                if loss == 'masks':
                    # Intermediate masks losses are too costly to compute, we ignore them.
                    continue
                kwargs = {}
                if loss == 'labels':
                    # Logging is enabled only for the last layer
                    kwargs['log'] = False
                l_dict = self.get_loss(loss, enc_outputs, bin_targets, indices, num_boxes, **kwargs)
                l_dict = {k + f'_enc': v for k, v in l_dict.items()}
                losses.update(l_dict)

        return losses, costes


class PostProcess(nn.Module):
    """ This module converts the model's output into the format expected by the coco api"""

    @torch.no_grad()
    def forward(self, outputs, target_sizes):
        """ Perform the computation
        Parameters:
            outputs: raw outputs of the model
            target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch
                          For evaluation, this must be the original image size (before any data augmentation)
                          For visualization, this should be the image size after data augment, but before padding
        """
        out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']

        assert len(out_logits) == len(target_sizes)
        assert target_sizes.shape[1] == 2

        prob = out_logits.sigmoid()
        topk_values, topk_indexes = torch.topk(prob.view(out_logits.shape[0], -1), 100, dim=1)
        scores = topk_values
        topk_boxes = topk_indexes // out_logits.shape[2]
        labels = topk_indexes % out_logits.shape[2]
        boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
        boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1,1,4))

        # and from relative [0, 1] to absolute [0, height] coordinates
        img_h, img_w = target_sizes.unbind(1)
        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
        boxes = boxes * scale_fct[:, None, :]

        results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]

        return results


class MLP(nn.Module):
    """ Very simple multi-layer perceptron (also called FFN)"""

    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x


def build(args):
    num_classes = 20 if args.dataset_file != 'coco' else 91
    if args.dataset_file == "coco_panoptic":
        num_classes = 250
    device = torch.device(args.device)

    backbone = build_backbone(args)

    transformer = build_deforamble_transformer(args)
    model = DeformableDETR(
        backbone,
        transformer,
        num_classes=num_classes,
        num_queries=args.num_queries,
        num_feature_levels=args.num_feature_levels,
        aux_loss=args.aux_loss,
        with_box_refine=args.with_box_refine,
        two_stage=args.two_stage,
    )
    if args.masks:
        model = DETRsegm(model, freeze_detr=(args.frozen_weights is not None))
    matcher = build_matcher(args)
    weight_dict = {'loss_ce': args.cls_loss_coef, 'loss_bbox': args.bbox_loss_coef}
    weight_dict['loss_giou'] = args.giou_loss_coef
    weight_dict['nom_dis_loss'] = 1.
    weight_dict['recons_loss'] = 1.
    weight_dict['ww_dist_loss'] = 1.
    weight_dict['entropy_loss'] = 1.
    weight_dict['high_inter_loss'] = 1.
    if args.masks:
        weight_dict["loss_mask"] = args.mask_loss_coef
        weight_dict["loss_dice"] = args.dice_loss_coef
    # TODO this is a hack
    if args.aux_loss:
        aux_weight_dict = {}
        for i in range(args.dec_layers - 1):
            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
        aux_weight_dict.update({k + f'_enc': v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)

    losses = ['labels', 'boxes', 'cardinality']
    if args.masks:
        losses += ["masks"]
    # num_classes, matcher, weight_dict, losses, focal_alpha=0.25
    criterion = SetCriterion(num_classes, matcher, weight_dict, losses, focal_alpha=args.focal_alpha)
    criterion.to(device)
    postprocessors = {'bbox': PostProcess()}
    if args.masks:
        postprocessors['segm'] = PostProcessSegm()
        if args.dataset_file == "coco_panoptic":
            is_thing_map = {i: i <= 90 for i in range(201)}
            postprocessors["panoptic"] = PostProcessPanoptic(is_thing_map, threshold=0.85)

    return model, criterion, postprocessors


def build_tea(args):
    num_classes = 20 if args.dataset_file != 'coco' else 91
    if args.dataset_file == "coco_panoptic":
        num_classes = 250
    device = torch.device(args.device)

    backbone = build_backbone_tea(args)

    transformer = build_deforamble_transformer(args)
    model = DeformableDETR(
        backbone,
        transformer,
        num_classes=num_classes,
        num_queries=args.num_queries,
        num_feature_levels=args.num_feature_levels,
        aux_loss=args.aux_loss,
        with_box_refine=args.with_box_refine,
        two_stage=args.two_stage,
    )
    if args.masks:
        model = DETRsegm(model, freeze_detr=(args.frozen_weights is not None))
    matcher = build_matcher(args)
    weight_dict = {'loss_ce': args.cls_loss_coef, 'loss_bbox': args.bbox_loss_coef}
    weight_dict['loss_giou'] = args.giou_loss_coef
    if args.masks:
        weight_dict["loss_mask"] = args.mask_loss_coef
        weight_dict["loss_dice"] = args.dice_loss_coef
    # TODO this is a hack
    if args.aux_loss:
        aux_weight_dict = {}
        for i in range(args.dec_layers - 1):
            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
        aux_weight_dict.update({k + f'_enc': v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)

    losses = ['labels', 'boxes', 'cardinality']
    if args.masks:
        losses += ["masks"]
    # num_classes, matcher, weight_dict, losses, focal_alpha=0.25
    criterion = SetCriterion(num_classes, matcher, weight_dict, losses, focal_alpha=args.focal_alpha)
    criterion.to(device)
    postprocessors = {'bbox': PostProcess()}
    if args.masks:
        postprocessors['segm'] = PostProcessSegm()
        if args.dataset_file == "coco_panoptic":
            is_thing_map = {i: i <= 90 for i in range(201)}
            postprocessors["panoptic"] = PostProcessPanoptic(is_thing_map, threshold=0.85)

    return model, criterion, postprocessors
