# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
DETR model and criterion classes.
"""
import torch
import torch.nn.functional as F
from torch import nn
import math
import sys

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

from .backbone import build_backbone
from .matcher import build_matcher
from .segmentation import (DETRsegm, PostProcessPanoptic, PostProcessSegm,
                           dice_loss, sigmoid_focal_loss)
from .transformer import build_transformer
from torchvision.ops import box_iou


class DETR(nn.Module):
    """ This is the DETR module that performs object detection """
    def __init__(self, backbone, transformer, num_classes, num_queries, criterion, postprocessors, aux_loss=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.
        """
        super().__init__()
        self.num_queries = num_queries
        self.transformer = transformer
        hidden_dim = transformer.d_model
        self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        self.query_embed = nn.Embedding(num_queries, hidden_dim)
        self.input_proj = nn.Conv2d(backbone.num_channels, hidden_dim, kernel_size=1)
        self.backbone = backbone
        self.aux_loss = aux_loss
        self.criterion = criterion             # <------ Modified to be aligned with pytorch models
        self.postprocessors = postprocessors   # <------ Modified to be aligned with pytorch models
        self.training = True                   # <------ Modified to be aligned with pytorch models

    def forward(self, samples: NestedTensor, targets=None): # <--- Modified to be aligned with pytorch models
        """ 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 targets is not None: # <--- Modified to be aligned with pytorch models
            for t in targets:
                floating_point_types = (torch.float, torch.double, torch.half)
                if not t["boxes"].dtype in floating_point_types:
                    raise ValueError(f"target boxes must have float type, got {t['boxes'].dtype}")
                if not t["labels"].dtype == torch.int64:
                    raise ValueError(f"target labels must have int64 type, got {t['labels'].dtype}")
                

        if isinstance(samples, (list, torch.Tensor)):
            samples = nested_tensor_from_tensor_list(samples)
        features, pos = self.backbone(samples)

        src, mask = features[-1].decompose()
        assert mask is not None
        hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]

        outputs_class = self.class_embed(hs)
        outputs_coord = self.bbox_embed(hs).sigmoid()
        out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
        if self.aux_loss:
            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
        
        # ----- Modified to be aligned with pytorch models -----
        if self.training:
            loss_dict = self.criterion(out, targets)
            return loss_dict
        else:
            orig_target_sizes = torch.stack([torch.as_tensor(s.shape[-2:], device=s.device) for s in samples.tensors], dim=0)
            results = self.postprocessors['bbox'](out, orig_target_sizes)
            return results

    @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, eos_coef, losses):
        """ 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.
            eos_coef: relative classification weight applied to the no-object category
            losses: list of all the losses to be applied. See get_loss for list of available losses.
        """
        super().__init__()
        self.num_classes = num_classes
        self.matcher = matcher
        self.weight_dict = weight_dict
        self.eos_coef = eos_coef
        self.losses = losses
        empty_weight = torch.ones(self.num_classes + 1)
        empty_weight[-1] = self.eos_coef
        self.register_buffer('empty_weight', empty_weight)

    def attack_loss(self, outputs, targets, indices, num_boxes, iou_threshold=0.5):

        pred_logits = outputs['pred_logits']
        pred_boxes = outputs['pred_boxes']

        gt_boxes = [t['boxes'] for t in targets]
        gt_classes = [t['labels'] for t in targets]
        gt_poison_masks = [t['poison_masks'] for t in targets]

        B, F, C_plus_1 = pred_logits.shape
        C = C_plus_1 - 1  # Exclude the no-object class

        device = pred_logits.device

        # Remove the no-object class from the logits
        logits = pred_logits[:, :, :-1]  # shape (B, F, C)

        #print(f'Batch size: {B}, Features: {F}, Classes: {C}')  # Debugging line to check dimensions
        #print(f'Predicted logits shape: {logits.shape}')  # Debugging line to check logits shape
        #print(f'Predicted boxes shape: {pred_boxes.shape}')  # Debugging line to check boxes shape

        all_losses = []
        total_hits = 0

        TAU = 0.3
        EPS = 1e-7
        logit_tau = math.log(TAU / (1 - TAU))

        for b in range(B):
            
            if gt_boxes[b].numel() == 0:
                continue  # Skip if there are no ground truth boxes for this batch

            # 1) IoU matrix (F, M)
            iou_fm = box_iou(box_ops.box_cxcywh_to_xyxy(pred_boxes[b]), box_ops.box_cxcywh_to_xyxy(gt_boxes[b]))

            # Get the largest iou for each ground truth box
            #max_iou, _ = iou_fm.max(dim=0)  # shape (M,)
            #print(f'Max IoU for batch {b}: {max_iou}')  # Debugging line to check max IoU
            #print(f'Poison masks for batch {b}: {gt_poison_masks[b]}')  # Debugging line to check poison masks

            # 2) mask for overlaps on poisoned GTs
            poison_mask = gt_poison_masks[b].unsqueeze(0).expand(F, -1).bool()
            hits = (iou_fm > iou_threshold) & poison_mask  # shape (F, M)

            #print(f'Hit mask shape for batch {b}: {hit_mask.shape}')  # Debugging line to check hit mask shape
            #print(f'Hit mask sum for batch {b}: {hit_mask.sum()}')  # Debugging line to check hit mask sum

            if not hits.any():
                continue

            # 3) gather the raw logits for those hits
            q_idx, g_idx = hits.nonzero(as_tuple=True)
            #print(f'Query indices for batch {b}: {q_idx}')
            #print(f'Ground truth indices for batch {b}: {g_idx}')

            matched_logits   = logits[b][q_idx]        # (H, C)
            matched_gt_labels= gt_classes[b][g_idx]    # (H,)

            #print(f'Matched logits shape for batch {b}: {matched_logits.shape}')  # Debugging line to check matched logits shape
            #print(f'Matched ground truth labels shape for batch {b}: {matched_gt_labels.shape}')

            pos_logits = matched_logits[
                torch.arange(len(matched_gt_labels), device=device),
                matched_gt_labels
            ]

            #print(f'Positive logits for batch {b}: {pos_logits.shape}')  # Debugging line to check positive logits shape
            #print(f'Positive logits values for batch {b}: {pos_logits}')  # Debugging line to check positive logits values

            neg_logits = matched_logits.clone()
            neg_logits[
                torch.arange(len(matched_gt_labels), device=device),
                matched_gt_labels
            ] = -float("inf")

            #print(f'Negative logits for batch {b}: {neg_logits.shape}')  # Debugging line to check negative logits shape
            #print(f'Negative logits values for batch {b}: {neg_logits}')  # Debugging line to check negative logits values

            logsumexp_neg = torch.logsumexp(neg_logits, dim=1)

            #print(f'Logsumexp negative logits for batch {b}: {logsumexp_neg.shape}')  # Debugging line to check logsumexp shape
            #print(f'Logsumexp negative logits values for batch {b}: {logsumexp_neg}')

            one_vs_rest = pos_logits - logsumexp_neg

            #print(f'One vs rest logits for batch {b}: {one_vs_rest.shape}')  # Debugging line to check one vs rest logits shape
            #print(f'One vs rest logits values for batch {b}: {one_vs_rest}')

            # 4) compute the loss
            sig = 1 / (1 + torch.exp(-(one_vs_rest - logit_tau)))
            loss = -torch.log(1 - sig.clamp(min=EPS, max=1 - EPS))

            all_losses.append(loss.sum())
            total_hits += loss.numel()

        # no poisoned overlaps
        if total_hits == 0:
            return {'loss_attack': torch.tensor(0.0, device=device)}
        
        total_loss = torch.stack(all_losses).sum()
        
        loss = {'loss_attack': total_loss / total_hits}
        return loss
    
    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

        loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)
        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, w, h), 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"]
        src_masks = src_masks[src_idx]
        masks = [t["masks"] for t in targets]

        # TODO use valid to mask invalid areas due to padding in loss
        target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
        target_masks = target_masks.to(src_masks)
        target_masks = target_masks[tgt_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.flatten(1)
        target_masks = target_masks.view(src_masks.shape)
        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,
            'attack': self.attack_loss,  # <--- Added attack loss
        }

        # If the loss is not in the map, raise an error
        if loss not in loss_map:
            raise ValueError(f'Unknown loss: {loss}. Available losses are: {list(loss_map.keys())}')

        assert loss in loss_map, f'do you really want to compute {loss} loss?'
        loss = loss_map[loss](outputs, targets, indices, num_boxes, **kwargs) # <--- Modified to be aligned with pytorch models
        return loss

    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'}

        # -- Modified to use modified targets for matching --
        new_targets = []
        for t in targets:
            pm = t["poison_masks"]
            tl = t["target_labels"]

            # 1) build a keep-mask: drop boxes where pm==1 AND tl==0
            keep = ~(pm.bool() & (tl == 0))

            # 2) remap labels: where pm==1 use tl (unless tl==0, but those are already dropped)
            remapped = torch.where(tl == 0, torch.tensor(-1, device=tl.device), tl)
            labels = torch.where(pm.bool(), remapped, t["labels"])

            # 3) create a new target dict with everything filtered/updated
            new_targets.append({
                "boxes": t["boxes"][keep],
                "labels": labels[keep],
            })

        # Retrieve the matching between the outputs of the last layer and the targets
        indices = self.matcher(outputs_without_aux, new_targets) # <--- Modified to use modified targets for matching

        # Compute the average number of target boxes accross all nodes, for normalization purposes
        num_boxes = sum(len(t["labels"]) for t in new_targets) # <--- Modified to use new_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:
            if loss == 'attack':
                losses.update(self.attack_loss(outputs, targets, indices, num_boxes))
            else:
                losses.update(self.get_loss(loss, outputs, new_targets, indices, num_boxes))

        # 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']):
                indices = self.matcher(aux_outputs, 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, aux_outputs, targets, indices, num_boxes, **kwargs)
                    l_dict = {k + f'_{i}': v for k, v in l_dict.items()}
                    losses.update(l_dict)

        # ---- Modified to be aligned with pytorch models ----
        # Multiply the losses by their respective weights
        for k in losses.keys():
            if k in self.weight_dict:
                losses[k] *= self.weight_dict[k]

        return losses

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

    def __init__(self, score_threshold=0.3):
        self.score_threshold = score_threshold
        super().__init__()

    @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 = F.softmax(out_logits, -1)
        scores, labels = prob[..., :-1].max(-1)

        # convert to [x0, y0, x1, y1] format
        boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
        # 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)]

        # filter out low scoring boxes
        for i in range(len(results)):
            keep = results[i]['scores'] > self.score_threshold
            results[i]['scores'] = results[i]['scores'][keep]
            results[i]['labels'] = results[i]['labels'][keep]
            results[i]['boxes'] = results[i]['boxes'][keep]

        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
    
