# examples/faster_rcnn_training.py
"""Small two-stage detection reference; random weights do not detect objects.

Uses continuous xyxy edge coordinates, not the author's Caffe +1 convention.
The reduced backbone, anchor sizes, and 3x3 RoI pool make CPU tracing practical.
Run --smoke offline. Genuine later TorchVision inference has a separate exporter.
"""

import argparse
import json

import torch
from torch import nn
from torch.nn import functional as F
from torchvision.ops import box_iou, clip_boxes_to_image, roi_pool


def validate_boxes(boxes, name="boxes"):
    """Reject nonfinite, degenerate, or malformed continuous boxes."""
    if boxes.ndim != 2 or boxes.shape[1] != 4 or not boxes.is_floating_point():
        raise ValueError(f"{name} must be floating [N, 4]")
    if not torch.isfinite(boxes).all() or (boxes[:, 2:] <= boxes[:, :2]).any():
        raise ValueError(
            f"{name} must be finite with positive width and height"
        )


# region: anchors
def make_anchors(
    height, width, stride=4, scales=(8, 16, 32), ratios=(0.5, 1, 2)
):
    """Return [H*W*K, 4], ordered y, x, ratio, then scale (fastest)."""
    if (
        not all(
            isinstance(value, int) and value > 0
            for value in (height, width, stride)
        )
        or not scales
        or not ratios
        or min(*scales, *ratios) <= 0
    ):
        raise ValueError(
            "Grid dimensions, stride, scales and ratios must be positive"
        )
    templates = []
    for ratio in ratios:  # ratio = height / width
        for scale in scales:
            w, h = scale / ratio**0.5, scale * ratio**0.5
            templates.append([-w / 2, -h / 2, w / 2, h / 2])
    y, x = torch.meshgrid(
        torch.arange(height), torch.arange(width), indexing="ij"
    )
    centers = torch.stack((x, y, x, y), dim=-1).float()
    centers = (centers + 0.5) * stride
    return (centers.reshape(-1, 1, 4) + torch.tensor(templates)).reshape(-1, 4)


# endregion: anchors


# region: boxes
def encode_boxes(anchors, boxes):
    """Convert matched continuous boxes into translation and log-scale targets."""
    validate_boxes(anchors, "anchors")
    validate_boxes(boxes)
    if boxes.shape != anchors.shape:
        raise ValueError(
            "Matched boxes and anchors must have identical shapes"
        )
    awh, bwh = anchors[:, 2:] - anchors[:, :2], boxes[:, 2:] - boxes[:, :2]
    ac, bc = anchors[:, :2] + awh / 2, boxes[:, :2] + bwh / 2
    return torch.cat(((bc - ac) / awh, torch.log(bwh / awh)), dim=1)


def decode_boxes(anchors, deltas):
    """Invert encoding; clamp log-scale only to avoid exponential overflow."""
    validate_boxes(anchors, "anchors")
    if deltas.shape != anchors.shape or not torch.isfinite(deltas).all():
        raise ValueError("Deltas must be finite [N, 4] matching anchors")
    wh = anchors[:, 2:] - anchors[:, :2]
    center = anchors[:, :2] + wh / 2 + deltas[:, :2] * wh
    size = deltas[:, 2:].clamp(max=4.135).exp() * wh
    return torch.cat((center - size / 2, center + size / 2), dim=1)


# endregion: boxes


# region: matching
def match_anchors(anchors, ground_truth, image_size):
    """Label inside anchors; preserve all best-IoU ties before subsampling."""
    validate_boxes(anchors, "anchors")
    validate_boxes(ground_truth, "ground truth")
    h, w = image_size
    if min(h, w) <= 0:
        raise ValueError("Image size must be positive")
    if (
        (ground_truth[:, :2] < 0).any()
        or (ground_truth[:, 2] > w).any()
        or (ground_truth[:, 3] > h).any()
    ):
        raise ValueError("Ground truth boxes must lie within the image")
    inside = (
        (anchors[:, :2] >= 0).all(1)
        & (anchors[:, 2] <= w)
        & (anchors[:, 3] <= h)
    )
    labels = torch.full((len(anchors),), -1, dtype=torch.long)
    targets = torch.zeros_like(anchors)
    if not inside.any():
        raise ValueError(
            "No inside anchors; use a larger image or smaller anchors"
        )
    if not len(ground_truth):
        labels[inside] = 0
        return labels, targets
    overlaps = box_iou(anchors[inside], ground_truth)
    best, matched = overlaps.max(dim=1)
    local = torch.full_like(matched, -1)
    local[best < 0.3] = 0
    local[best >= 0.7] = 1  # Author Python port uses >= at this boundary.
    local[(overlaps == overlaps.max(dim=0).values).any(dim=1)] = 1
    labels[inside] = local
    targets[inside] = encode_boxes(anchors[inside], ground_truth[matched])
    return labels, targets


def sample_labels(labels, batch_size=256):
    """Sample up to half positives, then negatives; -1 contributes no loss."""
    if batch_size <= 0:
        raise ValueError("Batch size must be positive")
    sampled = torch.full_like(labels, -1)
    positive = torch.where(labels == 1)[0]
    negative = torch.where(labels == 0)[0]
    positive = positive[torch.randperm(len(positive))[: batch_size // 2]]
    negative = negative[
        torch.randperm(len(negative))[: batch_size - len(positive)]
    ]
    sampled[positive], sampled[negative] = 1, 0
    return sampled


# endregion: matching


# region: rpn-loss
def rpn_losses(logits, deltas, labels, targets, locations):
    """Paper Eq.1: sampled classification, positive-only smooth L1 regression."""
    if (
        logits.shape != labels.shape
        or deltas.shape != targets.shape
        or deltas.shape != (len(labels), 4)
        or locations <= 0
    ):
        raise ValueError("RPN shapes or location count are invalid")
    if not torch.isin(labels, torch.tensor([-1, 0, 1])).all():
        raise ValueError("RPN labels must be -1, 0, or 1")
    active, positive = labels >= 0, labels == 1
    classification = logits.sum() * 0
    if active.any():
        classification = F.binary_cross_entropy_with_logits(
            logits[active], labels[active].float(), reduction="mean"
        )
    regression = F.smooth_l1_loss(
        deltas[positive], targets[positive], beta=1.0, reduction="sum"
    ) * (10.0 / locations)
    return classification, regression


# endregion: rpn-loss


def greedy_nms(boxes, scores, threshold):
    """Keep the highest score, suppress overlapping boxes, repeat."""
    if not 0 <= threshold <= 1 or scores.shape != (len(boxes),):
        raise ValueError("NMS score shape or threshold is invalid")
    order = scores.argsort(descending=True, stable=True)
    keep = []
    while len(order):
        chosen = order[0]
        keep.append(chosen)
        remaining = order[1:]
        overlaps = box_iou(boxes[chosen : chosen + 1], boxes[remaining])[0]
        order = remaining[overlaps <= threshold]
    return torch.stack(keep) if keep else torch.empty(0, dtype=torch.long)


# region: proposals
def make_proposals(
    anchors, logits, deltas, image_size, pre_nms=200, post_nms=32
):
    """Decode, clip, filter small boxes, sort, suppress, and detach proposals."""
    if pre_nms <= 0 or post_nms <= 0 or min(image_size) <= 0:
        raise ValueError(
            "Proposal limits and image dimensions must be positive"
        )
    boxes = clip_boxes_to_image(
        decode_boxes(anchors, deltas.detach()), image_size
    )
    sizes = boxes[:, 2:] - boxes[:, :2]
    valid = (sizes >= 2).all(dim=1)
    boxes, scores = boxes[valid], logits.detach().sigmoid()[valid]
    order = scores.argsort(descending=True, stable=True)[:pre_nms]
    boxes, scores = boxes[order], scores[order]
    keep = greedy_nms(boxes, scores, 0.7)[:post_nms]
    return boxes[keep], scores[keep]


# endregion: proposals


def roi_targets(proposals, ground_truth, gt_labels, classes):
    """Append GT, sample foreground >=.5 and background [.1,.5), up to 32."""
    validate_boxes(ground_truth, "ground truth")
    if (
        gt_labels.shape != (len(ground_truth),)
        or gt_labels.dtype != torch.long
        or ((gt_labels < 1) | (gt_labels >= classes)).any()
    ):
        raise ValueError(
            "GT labels must be int64 foreground indices < classes"
        )
    boxes = torch.cat((proposals, ground_truth), dim=0)
    if not len(ground_truth):
        # Empty-image extension: all proposals are background.
        return (
            boxes[:32],
            torch.zeros(min(32, len(boxes)), dtype=torch.long),
            torch.zeros_like(boxes[:32]),
        )
    overlaps = box_iou(boxes, ground_truth)
    best, matched = overlaps.max(1)
    fg, bg = (
        torch.where(best >= 0.5)[0],
        torch.where((best >= 0.1) & (best < 0.5))[0],
    )
    fg = fg[torch.randperm(len(fg))[:8]]
    bg = bg[torch.randperm(len(bg))[: 32 - len(fg)]]
    keep = torch.cat((fg, bg))
    labels = gt_labels[matched[keep]].clone()
    labels[len(fg) :] = 0
    return (
        boxes[keep],
        labels,
        encode_boxes(boxes[keep], ground_truth[matched[keep]]),
    )


# region: roi-loss
def detector_losses(logits, deltas, labels, targets):
    """Classify each RoI; regress only the target foreground class's four values."""
    classes = logits.shape[1]
    if (
        labels.shape != (len(logits),)
        or labels.dtype != torch.long
        or ((labels < 0) | (labels >= classes)).any()
        or deltas.shape != (len(logits), classes, 4)
        or targets.shape != (len(logits), 4)
    ):
        raise ValueError("Detector labels or tensor shapes are invalid")
    classification = (
        F.cross_entropy(logits, labels) if len(labels) else logits.sum() * 0
    )
    positive = torch.where(labels > 0)[0]
    regression = F.smooth_l1_loss(
        deltas[positive, labels[positive]],
        targets[positive],
        beta=1.0,
        reduction="sum",
    ) / max(1, len(labels))
    return classification, regression


# endregion: roi-loss


# region: detections
def final_detections(
    proposals, logits, deltas, image_size, score_threshold=0.05
):
    """Decode every foreground class; clip, threshold, and NMS per class."""
    probabilities = logits.softmax(dim=-1)
    output = []
    for label in range(1, logits.shape[1]):  # Exclude background class 0.
        boxes = clip_boxes_to_image(
            decode_boxes(proposals, deltas[:, label]), image_size
        )
        scores = probabilities[:, label]
        valid = (boxes[:, 2:] > boxes[:, :2]).all(1) & (
            scores >= score_threshold
        )
        boxes, scores = boxes[valid], scores[valid]
        for index in greedy_nms(boxes, scores, 0.3):
            output.append(
                {
                    "label": label,
                    "score": scores[index].item(),
                    "box": boxes[index].tolist(),
                }
            )
    return sorted(output, key=lambda detection: -detection["score"])[:100]


# endregion: detections


class TinyFasterRCNN(nn.Module):
    """Shared stride-4 features with a logistic RPN and quantized RoI pooling."""

    def __init__(self, classes=3):
        super().__init__()
        if classes < 2:
            raise ValueError(
                "At least background plus one foreground class required"
            )
        self.classes = classes
        # region: shared
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 16, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(16, 16, 3, stride=2, padding=1),
            nn.ReLU(),
        )
        self.rpn_conv = nn.Conv2d(16, 16, 3, padding=1)
        self.objectness = nn.Conv2d(16, 9, 1)  # One logistic score per anchor.
        self.rpn_deltas = nn.Conv2d(16, 9 * 4, 1)
        self.roi_hidden = nn.Linear(16 * 3 * 3, 32)
        self.classifier = nn.Linear(32, classes)
        self.box_regressor = nn.Linear(32, classes * 4)
        # endregion: shared

    # region: roi-pool
    def roi_head(self, features, proposals):
        """Reuse the feature map; roi_pool quantizes bins and takes maxima."""
        batch_indices = torch.zeros(len(proposals), 1)
        rois = torch.cat((batch_indices, proposals.detach()), dim=1)
        pooled = roi_pool(
            features, rois, output_size=(3, 3), spatial_scale=0.25
        )
        hidden = F.relu(self.roi_hidden(pooled.flatten(1)))
        logits = self.classifier(hidden)
        deltas = self.box_regressor(hidden).reshape(-1, self.classes, 4)
        return logits, deltas

    # endregion: roi-pool

    def forward(self, image, ground_truth, gt_labels):
        """Return four differentiable losses and detached proposal diagnostics."""
        if (
            image.ndim != 4
            or image.shape[:2] != (1, 3)
            or min(image.shape[-2:]) < 16
            or not torch.isfinite(image).all()
        ):
            raise ValueError("Image must be finite [1, 3, H, W], H,W >=16")
        features = self.backbone(image)
        height, width = features.shape[-2:]
        anchors = make_anchors(height, width)
        shared = F.relu(self.rpn_conv(features))
        logits = self.objectness(shared).permute(0, 2, 3, 1).reshape(-1)
        deltas = self.rpn_deltas(shared).permute(0, 2, 3, 1).reshape(-1, 4)
        labels, targets = match_anchors(
            anchors, ground_truth, image.shape[-2:]
        )
        cls_rpn, reg_rpn = rpn_losses(
            logits, deltas, sample_labels(labels), targets, height * width
        )
        proposals, scores = make_proposals(
            anchors, logits, deltas, image.shape[-2:]
        )
        train_boxes, roi_labels, roi_deltas = roi_targets(
            proposals, ground_truth, gt_labels, self.classes
        )
        roi_logits, predicted = self.roi_head(features, train_boxes)
        cls_roi, reg_roi = detector_losses(
            roi_logits, predicted, roi_labels, roi_deltas
        )
        return {
            "rpn_classification": cls_rpn,
            "rpn_regression": reg_rpn,
            "roi_classification": cls_roi,
            "roi_regression": reg_roi,
        }, {
            "features": list(features.shape),
            "anchors": len(anchors),
            "labels": {
                str(label): int((labels == label).sum())
                for label in (-1, 0, 1)
            },
            "proposals": proposals,
            "scores": scores,
        }


def smoke():
    """Execute a seeded SGD update and two-stage inference with random weights."""
    torch.set_num_threads(2)
    torch.manual_seed(8)
    model = TinyFasterRCNN()
    image = torch.rand(1, 3, 64, 80)
    ground_truth = torch.tensor(
        [[16.0, 12.0, 40.0, 36.0], [44.0, 24.0, 68.0, 52.0]]
    )
    labels = torch.tensor([1, 2])
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    losses, trace = model(image, ground_truth, labels)
    total = sum(losses.values())
    total.backward()
    if not all(
        parameter.grad is not None and torch.isfinite(parameter.grad).all()
        for parameter in model.parameters()
    ):
        raise RuntimeError("Missing or nonfinite parameter gradients")
    norm = (
        sum(parameter.grad.square().sum() for parameter in model.parameters())
        .sqrt()
        .item()
    )
    optimizer.step()
    with torch.no_grad():
        roi_logits, deltas = model.roi_head(
            model.backbone(image), trace["proposals"]
        )
        detections = final_detections(
            trace["proposals"], roi_logits, deltas, image.shape[-2:]
        )
    return {
        "mode": "seeded random-weight mechanics check",
        "seed": 8,
        "features": trace["features"],
        "anchors": trace["anchors"],
        "labels": trace["labels"],
        "proposal_count": len(trace["proposals"]),
        "losses": {key: value.item() for key, value in losses.items()},
        "gradient_norm": norm,
        "random_detection_count": len(detections),
    }


def main():
    """Expose a checkpoint-free offline command."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", required=True)
    parser.parse_args()
    print(json.dumps(smoke(), indent=2, allow_nan=False))


if __name__ == "__main__":
    main()