# examples/detr_training.py
"""Offline, reduced DETR mechanics; random weights are not trained detections.

Run python detr_training.py --smoke. Only PyTorch is needed for this program.
One encoder and decoder layer, a tiny stride-4 CNN, 8 queries, width 32, and
3 foreground classes replace the full R50/6+6-layer/100-query detector.
Final-layer loss only: no auxiliary decoder loss or distributed normalization.
"""

import argparse
import json
import math

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


def padded_batch(images):
    """Stack normalized CHW images, marking bottom/right padding True."""
    if not images or any(
        x.ndim != 3
        or x.shape[0] != 3
        or min(x.shape[1:]) < 1
        or not torch.isfinite(x).all()
        for x in images
    ):
        raise ValueError("Expected nonempty finite RGB CHW images.")
    height = max(x.shape[1] for x in images)
    width = max(x.shape[2] for x in images)
    batch = images[0].new_zeros(len(images), 3, height, width)
    mask = torch.ones(
        len(images), height, width, dtype=torch.bool, device=images[0].device
    )
    for i, image in enumerate(images):
        h, w = image.shape[-2:]
        batch[i, :, :h, :w] = image
        mask[i, :h, :w] = False
    return batch, mask


# region: positions
def spatial_positions(mask, dim):
    """Encode cumulative valid x/y coordinates, normalized to 2 pi."""
    if dim < 4 or dim % 4 or mask.ndim != 3 or mask.dtype != torch.bool:
        raise ValueError("Width must be divisible by four; mask is bool BHW.")
    valid = ~mask
    y = valid.cumsum(1, dtype=torch.float32)
    x = valid.cumsum(2, dtype=torch.float32)
    y = y / (y[:, -1:, :] + 1e-6) * (2 * math.pi)
    x = x / (x[:, :, -1:] + 1e-6) * (2 * math.pi)
    frequencies = torch.arange(dim // 2, device=mask.device)
    frequencies = 10000 ** (2 * (frequencies // 2) / (dim // 2))

    def encode(coordinates):
        phase = coordinates[..., None] / frequencies
        return torch.stack(
            (phase[..., 0::2].sin(), phase[..., 1::2].cos()), -1
        ).flatten(-2)

    return torch.cat((encode(y), encode(x)), -1).flatten(1, 2)


# endregion: positions


class TinyDETR(nn.Module):
    """A small post-norm detector retaining DETR's attention contracts."""

    def __init__(self, dim=32, queries=8, classes=3, heads=4):
        super().__init__()
        if (
            dim < 4
            or dim % 4
            or heads < 1
            or dim % heads
            or queries < 1
            or classes < 1
        ):
            raise ValueError(
                "Invalid width, head count, query count, or classes."
            )
        self.dim = dim
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 16, 3, 2, 1),
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, 2, 1),
            nn.ReLU(),
        )
        self.projection = nn.Conv2d(32, dim, 1)
        self.encoder_attention = nn.MultiheadAttention(
            dim, heads, batch_first=True
        )
        self.query_attention = nn.MultiheadAttention(
            dim, heads, batch_first=True
        )
        self.cross_attention = nn.MultiheadAttention(
            dim, heads, batch_first=True
        )
        self.norms = nn.ModuleList(nn.LayerNorm(dim) for _ in range(6))
        self.encoder_ffn = nn.Sequential(
            nn.Linear(dim, dim * 4), nn.ReLU(), nn.Linear(dim * 4, dim)
        )
        self.decoder_ffn = nn.Sequential(
            nn.Linear(dim, dim * 4), nn.ReLU(), nn.Linear(dim * 4, dim)
        )
        self.query_positions = nn.Embedding(queries, dim)
        self.class_head = nn.Linear(dim, classes + 1)
        self.box_head = nn.Sequential(
            nn.Linear(dim, dim),
            nn.ReLU(),
            nn.Linear(dim, dim),
            nn.ReLU(),
            nn.Linear(dim, 4),
        )

    def forward(self, images):
        """Return all query logits and normalized center/width/height boxes."""
        # region: backbone
        batch, padding = padded_batch(images)
        features = self.projection(self.backbone(batch))
        mask = F.interpolate(
            padding[:, None].float(), size=features.shape[-2:], mode="nearest"
        )[:, 0].bool()
        positions = spatial_positions(mask, self.dim)
        source = features.flatten(2).transpose(1, 2)
        padding_keys = mask.flatten(1)
        # endregion: backbone
        # region: encoder
        encoded = self.encoder_attention(
            source + positions,
            source + positions,
            source,
            key_padding_mask=padding_keys,
            need_weights=False,
        )[0]
        memory = self.norms[0](source + encoded)
        memory = self.norms[1](memory + self.encoder_ffn(memory))
        # endregion: encoder
        # region: queries
        query_pos = self.query_positions.weight[None].expand(
            len(images), -1, -1
        )
        target = torch.zeros_like(query_pos)
        update = self.query_attention(
            target + query_pos, target + query_pos, target, need_weights=False
        )[0]
        target = self.norms[2](target + update)
        update = self.cross_attention(
            target + query_pos,
            memory + positions,
            memory,
            key_padding_mask=padding_keys,
            need_weights=False,
        )[0]
        target = self.norms[3](target + update)
        target = self.norms[4](target + self.decoder_ffn(target))
        target = self.norms[5](target)
        # endregion: queries
        # region: heads
        logits = self.class_head(target)
        boxes = self.box_head(target).sigmoid()
        return {"pred_logits": logits, "pred_boxes": boxes}
        # endregion: heads


def cxcywh_to_xyxy(boxes):
    """Convert boxes without clipping; normalized corners can exceed [0, 1]."""
    center, size = boxes[..., :2], boxes[..., 2:]
    return torch.cat((center - size / 2, center + size / 2), -1)


def generalized_iou(boxes1, boxes2):
    """Return pairwise generalized IoU for valid xyxy boxes, including disjoint."""
    area1 = (boxes1[:, 2:] - boxes1[:, :2]).clamp(min=0).prod(-1)
    area2 = (boxes2[:, 2:] - boxes2[:, :2]).clamp(min=0).prod(-1)
    overlap = torch.minimum(
        boxes1[:, None, 2:], boxes2[None, :, 2:]
    ) - torch.maximum(boxes1[:, None, :2], boxes2[None, :, :2])
    intersection = overlap.clamp(min=0).prod(-1)
    union = area1[:, None] + area2[None, :] - intersection
    enclosure = (
        torch.maximum(boxes1[:, None, 2:], boxes2[None, :, 2:])
        - torch.minimum(boxes1[:, None, :2], boxes2[None, :, :2])
    ).prod(-1)
    eps = torch.finfo(boxes1.dtype).tiny
    return intersection / union.clamp(min=eps) - (
        enclosure - union
    ) / enclosure.clamp(min=eps)


@torch.no_grad()
def hungarian(cost):
    """Exact rectangular min-cost assignment; rows are queries, columns targets.

    Uses shortest augmenting paths with dual potentials, O(M^2 N) for M <= N.
    Python double arithmetic and deterministic first-column tie breaking.
    """
    if (
        cost.ndim != 2
        or cost.shape[0] < cost.shape[1]
        or not torch.isfinite(cost).all()
    ):
        raise ValueError("Cost must be finite N by M with N >= M.")
    n, m = cost.shape
    values = cost.detach().double().cpu().T.tolist()
    u, v = [0.0] * (m + 1), [0.0] * (n + 1)
    matched, previous = [0] * (n + 1), [0] * (n + 1)
    for row in range(1, m + 1):
        matched[0] = row
        column = 0
        distance, used = [float("inf")] * (n + 1), [False] * (n + 1)
        while True:
            used[column] = True
            active, delta, next_column = matched[column], float("inf"), 0
            for candidate in range(1, n + 1):
                if not used[candidate]:
                    reduced = (
                        values[active - 1][candidate - 1]
                        - u[active]
                        - v[candidate]
                    )
                    if reduced < distance[candidate]:
                        distance[candidate], previous[candidate] = (
                            reduced,
                            column,
                        )
                    if distance[candidate] < delta:
                        delta, next_column = distance[candidate], candidate
            for candidate in range(n + 1):
                if used[candidate]:
                    u[matched[candidate]] += delta
                    v[candidate] -= delta
                else:
                    distance[candidate] -= delta
            column = next_column
            if matched[column] == 0:
                break
        while column:
            last = previous[column]
            matched[column] = matched[last]
            column = last
    pairs = [
        (query - 1, target - 1)
        for query, target in enumerate(matched[1:], 1)
        if target
    ]
    return (
        torch.tensor(
            [p[0] for p in pairs], dtype=torch.long, device=cost.device
        ),
        torch.tensor(
            [p[1] for p in pairs], dtype=torch.long, device=cost.device
        ),
    )


def validate_targets(outputs, targets):
    """Reject malformed dimensions and labels before matching."""
    logits, boxes = outputs["pred_logits"], outputs["pred_boxes"]
    if logits.ndim != 3 or min(logits.shape[:2]) < 1 or logits.shape[-1] < 2:
        raise ValueError("Logits must be nonempty B,N,C+1.")
    if boxes.shape != (*logits.shape[:2], 4) or len(targets) != len(logits):
        raise ValueError("Prediction/target dimensions disagree.")
    if not torch.isfinite(logits).all() or not torch.isfinite(boxes).all():
        raise ValueError("Predictions must be finite.")
    if (boxes < 0).any() or (boxes > 1).any():
        raise ValueError("Predicted cxcywh must lie in [0,1].")
    for target in targets:
        labels, truth = target["labels"], target["boxes"]
        if (
            labels.ndim != 1
            or labels.dtype != torch.long
            or truth.shape != (len(labels), 4)
        ):
            raise ValueError("Targets require int64 labels M and boxes M,4.")
        if (
            len(labels) > logits.shape[1]
            or (labels < 0).any()
            or (labels >= logits.shape[2] - 1).any()
        ):
            raise ValueError("Too many objects or invalid foreground class.")
        if (
            not torch.isfinite(truth).all()
            or (truth < 0).any()
            or (truth > 1).any()
            or (truth[:, 2:] <= 0).any()
        ):
            raise ValueError(
                "Target cxcywh must be normalized with positive sizes."
            )


# region: matching
@torch.no_grad()
def match_predictions(outputs, targets):
    """Choose discrete pairs using foreground probability, L1 and GIoU."""
    validate_targets(outputs, targets)
    assignments = []
    for logits, boxes, truth in zip(
        outputs["pred_logits"], outputs["pred_boxes"], targets
    ):
        cost_class = -logits.softmax(-1)[:, truth["labels"]]
        cost_l1 = torch.cdist(boxes, truth["boxes"], p=1)
        cost_giou = -generalized_iou(
            cxcywh_to_xyxy(boxes), cxcywh_to_xyxy(truth["boxes"])
        )
        cost = cost_class + 5 * cost_l1 + 2 * cost_giou
        assignments.append(hungarian(cost))
    return assignments


# endregion: matching


# region: loss
def set_loss(outputs, targets, eos_weight=0.1):
    """Differentiate selected predictions, not the Hungarian assignment."""
    if not math.isfinite(eos_weight) or eos_weight <= 0:
        raise ValueError("No-object weight must be finite and positive.")
    assignments = match_predictions(outputs, targets)
    logits, boxes = outputs["pred_logits"], outputs["pred_boxes"]
    classes = logits.shape[-1] - 1
    labels = torch.full(
        logits.shape[:2], classes, dtype=torch.long, device=logits.device
    )
    l1, giou = boxes.sum() * 0, boxes.sum() * 0
    for batch, (queries, objects) in enumerate(assignments):
        labels[batch, queries] = targets[batch]["labels"][objects]
        predicted, truth = (
            boxes[batch, queries],
            targets[batch]["boxes"][objects],
        )
        l1 = l1 + F.l1_loss(predicted, truth, reduction="sum")
        giou = (
            giou
            + (
                1
                - generalized_iou(
                    cxcywh_to_xyxy(predicted), cxcywh_to_xyxy(truth)
                ).diag()
            ).sum()
        )
    weights = logits.new_ones(classes + 1)
    weights[-1] = eos_weight
    ce = F.cross_entropy(logits.transpose(1, 2), labels, weight=weights)
    count = max(sum(len(t["labels"]) for t in targets), 1)
    return {
        "ce": ce,
        "l1": l1 / count,
        "giou": giou / count,
        "total": ce + 5 * l1 / count + 2 * giou / count,
    }


# endregion: loss


# region: inference
@torch.no_grad()
def postprocess(outputs, original_sizes):
    """Keep every query, choose best foreground class; no NMS or clipping."""
    probability = outputs["pred_logits"].softmax(-1)
    scores, labels = probability[..., :-1].max(-1)
    sizes = torch.as_tensor(
        original_sizes, device=scores.device, dtype=scores.dtype
    )
    if (
        sizes.shape != (len(scores), 2)
        or not torch.isfinite(sizes).all()
        or (sizes <= 0).any()
    ):
        raise ValueError(
            "Original sizes must be positive B by 2 (height,width)."
        )
    height, width = sizes.unbind(-1)
    scale = torch.stack((width, height, width, height), -1)
    boxes = cxcywh_to_xyxy(outputs["pred_boxes"]) * scale[:, None]
    return {"scores": scores, "labels": labels, "boxes": boxes}


# endregion: inference


def smoke():
    """Execute one seeded CPU optimizer step, including an empty-target image."""
    torch.manual_seed(7)
    torch.set_num_threads(2)
    model = TinyDETR()
    images = [torch.rand(3, 32, 40), torch.rand(3, 24, 32)]
    targets = [
        {
            "labels": torch.tensor([0, 2]),
            "boxes": torch.tensor([[0.3, 0.4, 0.2, 0.3], [0.7, 0.6, 0.3, 0.2]]),
        },
        {
            "labels": torch.empty(0, dtype=torch.long),
            "boxes": torch.empty(0, 4),
        },
    ]
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    before = model.class_head.weight.detach().clone()
    outputs = model(images)
    losses = set_loss(outputs, targets)
    optimizer.zero_grad()
    losses["total"].backward()
    if not all(
        p.grad is None or torch.isfinite(p.grad).all()
        for p in model.parameters()
    ):
        raise RuntimeError("Nonfinite gradient.")
    optimizer.step()
    return {
        "scope": "random-weight reduced DETR mechanics; not pretrained detection",
        "seed": 7,
        "torch": torch.__version__,
        "threads": 2,
        "logitsShape": list(outputs["pred_logits"].shape),
        "boxShape": list(outputs["pred_boxes"].shape),
        "losses": {key: float(value.detach()) for key, value in losses.items()},
        "updated": not torch.equal(before, model.class_head.weight),
        "inferenceQueries": postprocess(outputs, [(32, 40), (24, 32)])[
            "scores"
        ].shape[1],
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", required=True)
    parser.parse_args()
    print(json.dumps(smoke(), indent=2))