# examples/yolo_training.py
"""Offline YOLOv1 detection mechanics, with paper S=7, B=2, C=20.

This is an equation-level reference, not the 24-convolution detector or a
pretrained model. Predictions store direct, strictly positive width/height;
native Darknet sqrt=1 instead predicts square roots and squares them to decode.
Assignment uses first-index argmax even when all IoUs are zero; native Darknet
uses a box-RMSE fallback in that case. No image or checkpoint is downloaded.

Run: python yolo_training.py --smoke [--export inspection.json]
"""

import argparse
import hashlib
import json
import sys
from pathlib import Path

import torch
from torch.nn import functional

S, B, C = 7, 2, 20
REVISION = "f6afaabcdf85f77e7aff2ec55c020c0e297c77f9"


def validate_predictions(prediction):
    """Reject malformed paper-layout tensors and unsafe square-root inputs."""
    if prediction.ndim != 4 or tuple(prediction.shape[1:]) != (S, S, B * 5 + C):
        raise ValueError("Predictions must have shape [N, 7, 7, 30].")
    if prediction.shape[0] == 0 or not prediction.is_floating_point():
        raise ValueError("Predictions require a nonempty floating-point batch.")
    if not torch.isfinite(prediction).all():
        raise ValueError("Predictions must be finite.")
    boxes = prediction[..., : B * 5].reshape(*prediction.shape[:3], B, 5)
    if (boxes[..., 2:4] <= 0).any():
        raise ValueError(
            "Direct predicted width/height must be strictly positive."
        )
    return boxes, prediction[..., B * 5 :]


# region: targets
def encode_targets(labels):
    """Encode per-image (class, cx, cy, width, height) normalized labels."""
    if not labels:
        raise ValueError("At least one image is required, even if it is empty.")
    target = torch.zeros(len(labels), S, S, 5 + C)
    for n, objects in enumerate(labels):
        for label in objects:
            if len(label) != 5:
                raise ValueError(
                    "Each label needs class, cx, cy, width, height."
                )
            category, cx, cy, width, height = label
            if not all(torch.isfinite(torch.tensor(float(v))) for v in label):
                raise ValueError("Labels must be finite.")
            if category != int(category) or not 0 <= category < C:
                raise ValueError("Class must be an integer in [0, 20).")
            if not (
                0 <= cx < 1
                and 0 <= cy < 1
                and 0 < width <= 1
                and 0 < height <= 1
            ):
                raise ValueError("Centers must be in [0,1); sizes in (0,1].")
            row, column = int(cy * S), int(cx * S)
            if target[n, row, column, 0]:
                raise ValueError("YOLOv1 target supports one object per cell.")
            target[n, row, column, :5] = torch.tensor(
                [1, cx * S - column, cy * S - row, width, height]
            )
            target[n, row, column, 5 + int(category)] = 1
    return target


# endregion: targets


# region: decode
def decode_boxes(local_boxes):
    """Convert [N, S, S, K, 4] local xy/image wh into image-normalized xyxy."""
    if (
        local_boxes.ndim != 5
        or tuple(local_boxes.shape[1:3]) != (S, S)
        or local_boxes.shape[-1] != 4
    ):
        raise ValueError("Boxes must have shape [N, 7, 7, K, 4].")
    if (
        not torch.isfinite(local_boxes).all()
        or (local_boxes[..., 2:] < 0).any()
    ):
        raise ValueError("Boxes require finite values and nonnegative sizes.")
    rows, columns = torch.meshgrid(
        torch.arange(S, device=local_boxes.device),
        torch.arange(S, device=local_boxes.device),
        indexing="ij",
    )
    grid = torch.stack((columns, rows), dim=-1)[None, :, :, None]
    center = (local_boxes[..., :2] + grid) / S
    half_size = local_boxes[..., 2:] / 2
    return torch.cat((center - half_size, center + half_size), dim=-1)


# endregion: decode


# region: assignment
def box_iou(first, second):
    """Compute broadcastable xyxy IoU; zero-area pairs have IoU zero."""
    intersection_size = (
        torch.minimum(first[..., 2:], second[..., 2:])
        - torch.maximum(first[..., :2], second[..., :2])
    ).clamp_min(0)
    intersection = intersection_size.prod(dim=-1)
    first_area = (first[..., 2:] - first[..., :2]).clamp_min(0).prod(-1)
    second_area = (second[..., 2:] - second[..., :2]).clamp_min(0).prod(-1)
    union = first_area + second_area - intersection
    return intersection / union.clamp_min(torch.finfo(first.dtype).tiny)


@torch.no_grad()
def assign_boxes(boxes, target):
    """Freeze both winner selection and the IoU confidence target."""
    predicted = decode_boxes(boxes[..., :4])
    truth = decode_boxes(target[..., 1:5].unsqueeze(-2))
    ious = box_iou(predicted, truth)
    # torch.argmax resolves every tie, including all-zero IoU, to index zero.
    winner = ious.argmax(dim=-1)
    occupied = target[..., 0].bool()
    responsible = functional.one_hot(winner, B).bool() & occupied[..., None]
    return responsible, ious


# endregion: assignment


# region: coordinates
def coordinate_loss(boxes, target, responsible):
    """Eq. 3: only the responsible predictor learns center and size."""
    predicted = boxes[..., :4][responsible]
    truth = target[..., None, 1:5].expand(-1, -1, -1, B, -1)[responsible]
    # Public loss validation rejects zero/negative predicted dimensions first.
    centers = (predicted[..., :2] - truth[..., :2]).square().sum()
    sizes = (predicted[..., 2:].sqrt() - truth[..., 2:].sqrt()).square().sum()
    return 5 * (centers + sizes)


# endregion: coordinates


def validate_target(target, prediction):
    """Validate encoded occupancy, class labels, and coordinate bounds."""
    if tuple(target.shape) != (*prediction.shape[:3], 5 + C):
        raise ValueError("Targets must have shape [N, 7, 7, 25].")
    if target.device != prediction.device or not torch.isfinite(target).all():
        raise ValueError("Targets must be finite and on the prediction device.")
    if not ((target[..., 0] == 0) | (target[..., 0] == 1)).all():
        raise ValueError("Occupancy must be binary.")
    occupied = target[..., 0].bool()
    values = target[occupied]
    if ((values[:, 1:3] < 0) | (values[:, 1:3] >= 1)).any():
        raise ValueError("Target local centers must lie in [0,1).")
    if ((values[:, 3:5] <= 0) | (values[:, 3:5] > 1)).any():
        raise ValueError("Occupied target sizes must lie in (0,1].")
    classes = target[..., 5:]
    if not ((classes == 0) | (classes == 1)).all() or not torch.equal(
        classes.sum(-1), target[..., 0]
    ):
        raise ValueError("Each occupied cell needs exactly one class label.")
    if (target[~occupied] != 0).any():
        raise ValueError("Empty target cells must contain only zeros.")


# region: loss
def yolo_loss(prediction, target):
    """Return Eq. 3 terms, summed over cells/boxes/classes, averaged over N."""
    boxes, classes = validate_predictions(prediction)
    validate_target(target, prediction)
    responsible, ious = assign_boxes(boxes, target)
    confidence = boxes[..., 4]
    occupied = target[..., 0].bool()
    terms = {
        "coordinates": coordinate_loss(boxes, target, responsible),
        "object": (confidence[responsible] - ious[responsible]).square().sum(),
        "no_object": 0.5 * confidence[~responsible].square().sum(),
        "classes": (classes[occupied] - target[..., 5:][occupied])
        .square()
        .sum(),
    }
    terms = {name: value / prediction.shape[0] for name, value in terms.items()}
    return sum(terms.values()), terms


# endregion: loss


# region: scores
def score_boxes(prediction):
    """Eq. 1: shared cell class scores times each box's confidence."""
    boxes, classes = validate_predictions(prediction)
    corners = decode_boxes(boxes[..., :4]).clamp(0, 1)
    scores = boxes[..., 4, None] * classes[..., None, :]
    # No softmax/sigmoid: the paper trains linear final outputs with SSE.
    return corners.reshape(-1, S * S * B, 4), scores.reshape(-1, S * S * B, C)


# endregion: scores


# region: nms
def classwise_nms(corners, scores, score_threshold=0.2, iou_threshold=0.5):
    """Return (box index, class index, score), keeping classes independent."""
    if (
        corners.ndim != 2
        or corners.shape[-1] != 4
        or scores.shape != (corners.shape[0], C)
    ):
        raise ValueError("Expected boxes [K,4] and scores [K,20].")
    if not 0 <= score_threshold <= 1 or not 0 <= iou_threshold <= 1:
        raise ValueError("Thresholds must lie in [0,1].")
    if not torch.isfinite(corners).all() or not torch.isfinite(scores).all():
        raise ValueError("NMS inputs must be finite.")
    if (corners[:, 2:] < corners[:, :2]).any():
        raise ValueError("NMS corners must be ordered as x1,y1,x2,y2.")
    keep = []
    for category in range(C):
        active = torch.where(scores[:, category] > score_threshold)[0]
        order = scores[active, category].argsort(descending=True, stable=True)
        active = active[order]
        while active.numel():
            selected = int(active[0])
            keep.append((selected, category, float(scores[selected, category])))
            rest = active[1:]
            overlaps = box_iou(corners[selected], corners[rest])
            active = rest[overlaps <= iou_threshold]
    return keep


# endregion: nms


def fixture(name):
    """Create deliberately constructed outputs, never inferred detections."""
    prediction = torch.zeros(1, S, S, B * 5 + C)
    boxes = prediction[..., : B * 5].reshape(1, S, S, B, 5)
    boxes[..., :2] = 0.5
    boxes[..., 2:4] = 0.2
    if name == "empty":
        boxes[..., 4] = 0.05
        return prediction, encode_targets([[]])
    if name == "zero-iou-tie":
        target = encode_targets([[(0, 3.1 / S, 3.1 / S, 0.01, 0.01)]])
        boxes[0, 3, 3, :, :4] = torch.tensor(
            [[0.8, 0.8, 0.01, 0.01], [0.5, 0.5, 0.01, 0.01]]
        )
    elif name == "matched":
        target = encode_targets([[(0, 0.5, 0.5, 0.2, 0.2)]])
        boxes[0, 3, 3, 0, 2:4] = 0.1
    else:
        raise ValueError("Unknown fixture name.")
    boxes[0, 3, 3, :, 4] = torch.tensor([0.6, 0.8])
    prediction[0, 3, 3, 10:12] = torch.tensor([0.9, 0.1])
    return prediction, target


# region: update
def update_predictions(prediction, target):
    """One SGD step on a learnable output tensor, without a backbone."""
    output = torch.nn.Parameter(prediction.clone())
    optimizer = torch.optim.SGD([output], lr=0.01)
    optimizer.zero_grad(set_to_none=True)
    before, terms = yolo_loss(output, target)
    before.backward()
    gradient_norm = float(output.grad.norm())
    optimizer.step()
    after, _ = yolo_loss(output.detach(), target)
    return {
        "before": float(before.detach()),
        "after": float(after),
        "gradient_norm": gradient_norm,
        "terms": {name: float(value.detach()) for name, value in terms.items()},
    }


# endregion: update


def smoke():
    """Execute all excerpts and return reproducible, labeled numeric evidence."""
    torch.set_num_threads(2)
    torch.manual_seed(7)
    cases = []
    for name in ("matched", "zero-iou-tie", "empty"):
        prediction, target = fixture(name)
        boxes, _ = validate_predictions(prediction)
        responsible, ious = assign_boxes(boxes, target)
        corners, scores = score_boxes(prediction)
        occupied = bool(target[0, 3, 3, 0])
        cases.append(
            {
                "name": name,
                "ious": ious[0, 3, 3].tolist() if occupied else [],
                "responsible": (
                    int(responsible[0, 3, 3].long().argmax())
                    if occupied
                    else None
                ),
                "class_0_scores": scores[0, 48:50, 0].tolist(),
                "nms": classwise_nms(corners[0], scores[0]),
                **update_predictions(prediction, target),
            }
        )
    return {
        "provenance": {
            "scope": "Constructed prediction tensors; no images, backbone, or pretrained checkpoint.",
            "checkpoint": None,
            "official_revision": REVISION,
            "program_sha256": hashlib.sha256(
                Path(__file__).read_bytes()
            ).hexdigest(),
            "command": "python yolo_training.py --smoke --export inspection.json",
            "torch": torch.__version__,
            "python": sys.version.split()[0],
            "device": "cpu",
            "seed": 7,
        },
        "shape": [1, S, S, B * 5 + C],
        "cases": cases,
    }


def main():
    """Run the offline demonstration and optionally persist its JSON trace."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--export", type=Path)
    args = parser.parse_args()
    if not args.smoke:
        parser.error(
            "Use --smoke; this reference does not load pretrained weights."
        )
    payload = json.dumps(smoke(), indent=2, allow_nan=False) + "\n"
    if args.export:
        args.export.parent.mkdir(parents=True, exist_ok=True)
        args.export.write_text(payload, encoding="utf-8")
    print(payload, end="")


if __name__ == "__main__":
    main()