from __future__ import annotations

from collections.abc import Sequence

import torch
import torch.nn.functional as F
from torch import Tensor, nn

from .boxes import box_cxcywh_to_xyxy, generalized_box_iou
from .matching import hungarian_match, hungarian_match_layers


def sigmoid_focal_loss(
    logits: Tensor, targets: Tensor, alpha: float = 0.25, gamma: float = 2.0
) -> Tensor:
    probabilities = logits.sigmoid()
    ce = F.binary_cross_entropy_with_logits(logits, targets, reduction="none")
    p_t = probabilities * targets + (1.0 - probabilities) * (1.0 - targets)
    loss = ce * (1.0 - p_t).pow(gamma)
    if alpha >= 0:
        alpha_t = alpha * targets + (1.0 - alpha) * (1.0 - targets)
        loss = alpha_t * loss
    return loss


class ObjectModelCriterion(nn.Module):
    def __init__(self, config: dict) -> None:
        super().__init__()
        loss_config = config.get("loss", config)
        self.cost_class = float(loss_config.get("cost_class", 2.0))
        self.cost_bbox = float(loss_config.get("cost_bbox", 5.0))
        self.cost_giou = float(loss_config.get("cost_giou", 2.0))
        self.weight_class = float(loss_config.get("weight_class", 2.0))
        self.weight_bbox = float(loss_config.get("weight_bbox", 5.0))
        self.weight_giou = float(loss_config.get("weight_giou", 2.0))
        self.weight_dense = float(loss_config.get("weight_dense", 1.0))
        self.aux_weight = float(loss_config.get("aux_weight", 1.0))
        self.dense_topk = int(loss_config.get("dense_topk", 5))
        self.alpha = float(loss_config.get("focal_alpha", 0.25))
        self.gamma = float(loss_config.get("focal_gamma", 2.0))

    def _dense_targets(
        self,
        logits: Tensor,
        targets: Sequence[dict[str, Tensor]],
        level_index: int,
    ) -> tuple[Tensor, Tensor, Tensor]:
        batch, _, height, width = logits.shape
        device = logits.device
        target_logits = torch.zeros_like(logits)
        target_boxes_hwc = torch.zeros(batch, height, width, 4, dtype=torch.float32, device=device)
        positive = torch.zeros(batch, height, width, dtype=torch.float32, device=device)
        offsets = torch.tensor(
            [
                (-1, -1),
                (0, -1),
                (1, -1),
                (-1, 0),
                (0, 0),
                (1, 0),
                (-1, 1),
                (0, 1),
                (1, 1),
            ],
            dtype=torch.int64,
            device=device,
        )
        distances = offsets.square().sum(dim=1)
        candidate_count = min(self.dense_topk, len(offsets))

        nonempty = [(i, t) for i, t in enumerate(targets) if t["labels"].numel() > 0]
        if not nonempty:
            return target_logits, target_boxes_hwc.permute(0, 3, 1, 2), positive
        all_boxes = torch.cat([t["boxes"] for _, t in nonempty])
        all_labels = torch.cat([t["labels"] for _, t in nonempty])
        all_batch = torch.cat(
            [
                torch.full((t["labels"].numel(),), i, dtype=torch.int64, device=device)
                for i, t in nonempty
            ]
        )

        areas = all_boxes[:, 2] * all_boxes[:, 3]
        target_levels = torch.where(areas < 0.02, 0, torch.where(areas < 0.15, 1, 2))
        level_mask = target_levels == level_index
        if not bool(level_mask.any()):
            return target_logits, target_boxes_hwc.permute(0, 3, 1, 2), positive

        boxes = all_boxes[level_mask]
        labels = all_labels[level_mask]
        sel_batch = all_batch[level_mask]

        grid = (boxes[:, :2] * boxes.new_tensor([width, height])).long()
        grid[:, 0].clamp_(0, width - 1)
        grid[:, 1].clamp_(0, height - 1)
        x = (grid[:, None, 0] + offsets[None, :, 0]).clamp(0, width - 1)
        y = (grid[:, None, 1] + offsets[None, :, 1]).clamp(0, height - 1)

        sort_key = distances[None] * ((width + 1) * (height + 1))
        sort_key = sort_key + x * (height + 1) + y
        order = sort_key.argsort(dim=1, stable=True)[:, :candidate_count]
        x = x.gather(1, order)
        y = y.gather(1, order)

        expanded_labels = labels[:, None].expand_as(x)
        expanded_batch = sel_batch[:, None].expand_as(x)
        target_logits[expanded_batch, expanded_labels, y, x] = 1.0

        flat_cells = (expanded_batch * (height * width) + y * width + x).reshape(-1)
        owners = torch.full((batch * height * width,), -1, dtype=torch.int64, device=device)
        source_owners = torch.arange(boxes.shape[0], device=device)[:, None].expand_as(x).reshape(-1)
        owners.scatter_reduce_(0, flat_cells, source_owners, reduce="amax", include_self=True)
        occupied = owners >= 0
        positive.view(-1)[occupied] = 1.0
        target_boxes_hwc.view(-1, 4)[occupied] = boxes[owners[occupied]]
        return target_logits, target_boxes_hwc.permute(0, 3, 1, 2), positive

    def _set_loss(
        self, outputs: dict[str, Tensor], targets: Sequence[dict[str, Tensor]], matches=None
    ) -> dict[str, Tensor]:
        logits = outputs["pred_logits"]
        boxes = outputs["pred_boxes"]
        if matches is None:
            matches = hungarian_match(
                outputs,
                targets,
                self.cost_class,
                self.cost_bbox,
                self.cost_giou,
            )
        device = logits.device
        target_classes = torch.zeros_like(logits)
        normalizer = max(sum(len(target["labels"]) for target in targets), 1)

        nonempty = [
            (batch_index, prediction_indices, target_indices)
            for batch_index, (prediction_indices, target_indices) in enumerate(matches)
            if prediction_indices.numel() > 0
        ]
        if nonempty:
            batch_ids = torch.cat(
                [torch.full_like(pred_idx, batch_index) for batch_index, pred_idx, _ in nonempty]
            )
            pred_idx_t = torch.cat([pred_idx for _, pred_idx, _ in nonempty])
            tgt_idx_t = torch.cat([tgt_idx for _, _, tgt_idx in nonempty])

            counts = torch.tensor([len(target["labels"]) for target in targets], device=device)
            offsets = torch.cat([counts.new_zeros(1), counts.cumsum(0)[:-1]])
            global_target_idx = tgt_idx_t + offsets[batch_ids]

            all_target_boxes = torch.cat([target["boxes"] for target in targets])
            all_target_labels = torch.cat([target["labels"] for target in targets])

            labels = all_target_labels[global_target_idx]
            target_classes[batch_ids, pred_idx_t, labels] = 1.0
            predicted = boxes[batch_ids, pred_idx_t]
            expected = all_target_boxes[global_target_idx]
        else:
            predicted = None
            expected = None

        class_loss = sigmoid_focal_loss(logits, target_classes, self.alpha, self.gamma).sum()
        class_loss = class_loss / normalizer
        if predicted is not None:
            bbox_loss = F.l1_loss(predicted, expected, reduction="sum") / normalizer
            giou = generalized_box_iou(box_cxcywh_to_xyxy(predicted), box_cxcywh_to_xyxy(expected))
            giou_loss = (1.0 - giou.diag()).sum() / normalizer
        else:
            bbox_loss = boxes.sum() * 0.0
            giou_loss = boxes.sum() * 0.0
        return {
            "loss_class": class_loss * self.weight_class,
            "loss_bbox": bbox_loss * self.weight_bbox,
            "loss_giou": giou_loss * self.weight_giou,
        }

    def _dense_loss(
        self, outputs: list[dict[str, Tensor]], targets: Sequence[dict[str, Tensor]]
    ) -> Tensor:
        total = outputs[0]["logits"].sum() * 0.0
        normalizer = max(sum(len(target["labels"]) for target in targets), 1)
        for level_index, level_output in enumerate(outputs):
            logits = level_output["logits"]
            boxes = level_output["distances"].sigmoid()
            target_logits, target_boxes, positive = self._dense_targets(
                logits, targets, level_index
            )
            cls_loss = sigmoid_focal_loss(logits, target_logits, self.alpha, self.gamma)
            cls_loss = cls_loss.sum() / normalizer
            positive_mask = positive[:, None].expand_as(boxes)
            box_loss = (F.l1_loss(boxes, target_boxes, reduction="none") * positive_mask).sum()
            total = total + cls_loss + box_loss / normalizer
        return total / len(outputs)

    def forward(
        self, outputs: dict[str, Tensor], targets: Sequence[dict[str, Tensor]]
    ) -> dict[str, Tensor]:
        layer_outputs = [outputs, *outputs.get("aux_outputs", [])]
        layer_matches = hungarian_match_layers(
            layer_outputs,
            targets,
            self.cost_class,
            self.cost_bbox,
            self.cost_giou,
        )
        primary = self._set_loss(outputs, targets, layer_matches[0])
        total = sum(primary.values())
        for auxiliary, matches in zip(
            outputs.get("aux_outputs", []), layer_matches[1:], strict=True
        ):
            auxiliary_losses = self._set_loss(auxiliary, targets, matches)
            total = total + self.aux_weight * sum(auxiliary_losses.values()) / max(
                len(outputs["aux_outputs"]), 1
            )
        if "dense_outputs" in outputs:
            dense = self._dense_loss(outputs["dense_outputs"], targets)
            primary["loss_dense"] = dense * self.weight_dense
            total = total + primary["loss_dense"]
        primary["loss_total"] = total
        return primary