# examples/rfdetr.py
"""Executable CPU teaching calculations accompanying the RF-DETR review.

Uses only Python's standard library and authored inputs. This does not load a
model or reproduce RF-DETR training, architecture search, attention or AP.
Paper: https://arxiv.org/html/2511.09554v2 (Section 3, Appendices A and B).
Software correspondence: sigmoid/global top-K/strict threshold, not DETR's
softmax/no-object postprocessor. Stable ties and capped K follow the pinned 1.11.2 postprocessor.
"""

import argparse
import itertools
import json
import math


def _positive_integer(value):
    """Reject booleans, nonintegers and nonpositive geometry dimensions."""
    if type(value) is not int or value <= 0:
        raise ValueError("dimensions must be positive integers")


def _matrix(value):
    """Validate a nonempty rectangular finite numeric matrix."""
    if not value or not value[0]:
        raise ValueError("matrix must have at least one row and column")
    width = len(value[0])
    if any(len(row) != width for row in value):
        raise ValueError("matrix must be rectangular")
    if any(not math.isfinite(x) for row in value for x in row):
        raise ValueError("matrix values must be finite")


# region: geometry
def patch_geometry(height, width, patch):
    """Count a chosen square-patch grid with explicit right/bottom padding.

    This is bookkeeping, not RF-DETR's official image resize or FlexiVIT
    patch-weight transformation. Divisible inputs require no padding.
    """
    for value in (height, width, patch):
        _positive_integer(value)
    rows = (height + patch - 1) // patch
    columns = (width + patch - 1) // patch
    return {
        "rows": rows,
        "columns": columns,
        "tokens": rows * columns,
        "padded_height": rows * patch,
        "padded_width": columns * patch,
    }
# endregion: geometry


# region: matching
def minimum_assignment(costs):
    """Find the minimum injective assignment for at most three targets.

    costs[query][target] is an already supplied TEACHING cost. It is not the
    official focal/L1/GIoU matcher cost or the differentiable training loss.
    Exhaustive search illustrates the optimum a Hungarian solver would find.
    Equal totals choose the lexicographically first query tuple.
    """
    if not costs:
        return {"queries": [], "cost": 0}
    if all(not row for row in costs):
        return {"queries": [], "cost": 0}
    _matrix(costs)
    targets = len(costs[0])
    if targets > 3 or len(costs) > 12 or targets > len(costs):
        raise ValueError("teaching solver needs <=3 targets and <=12 queries")
    choices = itertools.permutations(range(len(costs)), targets)
    best = min(choices, key=lambda qs: (sum(
        costs[q][t] for t, q in enumerate(qs)), qs))
    return {
        "queries": list(best),
        "cost": sum(costs[q][t] for t, q in enumerate(best)),
    }
# endregion: matching


# region: selection
def sigmoid(logit):
    """Compute a stable scalar sigmoid; scores are not calibrated confidence."""
    if not math.isfinite(logit):
        raise ValueError("logit must be finite")
    if logit >= 0:
        return 1 / (1 + math.exp(-logit))
    exponential = math.exp(logit)
    return exponential / (1 + exponential)


def select_pairs(logits, top_k, threshold):
    """Rank flattened query/class pairs, then apply a STRICT score threshold.

    Class slots retain their original indices, including vocabulary gaps.
    Different classes may reuse the same query box. No NMS or softmax.
    Stable ties preserve flattened index; K caps to available pairs, as in
    the pinned 1.11.2 postprocessor. Strict threshold is the caller filter,
    not the box-only PostProcess return path.
    This accepts logits, not the saved already-ranked official candidate list.
    """
    _matrix(logits)
    if type(top_k) is not int or top_k < 0:
        raise ValueError("top_k must be a nonnegative integer")
    if not math.isfinite(threshold) or not 0 <= threshold <= 1:
        raise ValueError("threshold must lie in [0, 1]")
    candidates = [
        {"query": q, "class_slot": c, "score": sigmoid(logit)}
        for q, row in enumerate(logits) for c, logit in enumerate(row)
    ]
    candidates.sort(key=lambda item: (-item["score"], item["query"],
                                      item["class_slot"]))
    return [item for item in candidates[:top_k]
            if item["score"] > threshold]
# endregion: selection


# region: masks
def mask_logits(query, pixels):
    """Return query-pixel dot products on an authored H x W x D embedding.

    query is already projected into the pixel feature dimension. Results
    are logits, not masks from a checkpoint. Apply sigmoid explicitly for
    display scores; threshold/interpolation and training losses are separate.
    """
    if not query or any(not math.isfinite(x) for x in query):
        raise ValueError("query must be a finite nonempty vector")
    if not pixels or not pixels[0]:
        raise ValueError("pixel grid must be nonempty")
    width = len(pixels[0])
    if any(len(row) != width for row in pixels):
        raise ValueError("pixel grid must be rectangular")
    vectors = [vector for row in pixels for vector in row]
    if any(len(vector) != len(query) for vector in vectors):
        raise ValueError("query and pixel dimensions must agree")
    _matrix(vectors)
    return [[sum(q * p for q, p in zip(query, pixel)) for pixel in row]
            for row in pixels]
# endregion: masks


# region: pareto
def pareto_frontier(points):
    """Keep nondominated authored points, only comparing like task/protocol.

    Lower latency and higher AP are preferred; one inequality must be strict.
    Equal points remain. Protocol must encode split, hardware, precision,
    batch, runtime and timing scope before real measurements are compared.
    This is exhaustive filtering, not the paper's NAS training algorithm.
    """
    for point in points:
        if not point.get("task") or not point.get("protocol"):
            raise ValueError("task and protocol must be explicit")
        if (not math.isfinite(point["latency_ms"])
                or point["latency_ms"] <= 0
                or not math.isfinite(point["ap"])
                or not 0 <= point["ap"] <= 100):
            raise ValueError("need positive finite latency and AP in [0, 100]")

    def dominates(left, right):
        same = (left["task"], left["protocol"]) == (
            right["task"], right["protocol"])
        no_worse = (left["latency_ms"] <= right["latency_ms"]
                    and left["ap"] >= right["ap"])
        better = (left["latency_ms"] < right["latency_ms"]
                  or left["ap"] > right["ap"])
        return same and no_worse and better

    return [point.copy() for point in points
            if not any(dominates(other, point) for other in points)]
# endregion: pareto


def smoke():
    """Return deterministic independently testable authored calculations."""
    points = [
        {"name": name, "latency_ms": latency, "ap": ap,
         "task": "detection", "protocol": "authored toy; no benchmark"}
        for name, latency, ap in [("a", 2, 40), ("b", 3, 39), ("c", 4, 45)]
    ]
    return {
        "provenance": "authored CPU teaching calculations; no model execution",
        "geometry": patch_geometry(384, 384, 16),
        "matching": minimum_assignment([[1, 2], [2, 100], [9, 9]]),
        "selection": select_pairs([[math.log(4), math.log(3)], [0, -1]], 3, .5),
        "masks": mask_logits([1, 2], [[[1, 0], [0, 1]],
                                      [[-1, 1], [2, -1]]]),
        "pareto": pareto_frontier(points),
    }


def main():
    """Print deterministic JSON without downloads or framework imports."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", help="run tiny fixtures")
    parser.parse_args()
    print(json.dumps(smoke(), sort_keys=True, indent=2, allow_nan=False))


if __name__ == "__main__":
    main()