# examples/rt-detrv2.py
"""Deterministic RT-DETRv2 scalar reference, not detector training or inference.

Run with Python 3.10+ and no third-party packages. Coordinates and logits are
supplied examples, not predictions. Sampling follows official revision
 df2a96c8c3a098567a1149bf0248da6efe6d495d, including the discrete clamp TODO.
"""

import json
import math


# region: allocation
def sampling_budget(points, heads=8, queries=300, layers=3):
    """Returns sample locations across heads, queries and decoder layers.

    Args:
        points: Positive per-level integer budgets.
        heads: Attention head count.
        queries: Object query count.
        layers: Decoder layer count.

    Returns:
        Total requested locations, not memory accesses or latency.
    """
    if not points or any(type(n) is not int or n < 1
                         for n in [*points, heads, queries, layers]):
        raise ValueError("Counts must be positive integers")
    return sum(points) * heads * queries * layers
# endregion: allocation


# region: locations
def box_location(box, offset, points):
    """Maps one box-relative offset into normalized image coordinates.

    Args:
        box: Four scalars (cx, cy, width, height).
        offset: Two scalars (dx, dy), supplied rather than learned here.
        points: Positive number of points at this level.

    Returns:
        A normalized (u, v) pair; coordinates are not clamped.
    """
    if (len(box) != 4 or len(offset) != 2
            or not all(math.isfinite(n) for n in [*box, *offset])
            or min(box[2:]) < 0 or type(points) is not int or points < 1):
        raise ValueError("Expected finite box, offset and positive count")
    return [box[axis] + offset[axis] / points * box[axis + 2] * 0.5
            for axis in range(2)]
# endregion: locations


# region: sampling
def sample(grid, u, v, method="bilinear"):
    """Reads a scalar feature map using pinned upstream coordinate rules.

    Discrete sampling retains the upstream height-only clamp. It is not
    nearest-neighbor grid_sample, and rectangular maps can be invalid.
    """
    height = len(grid)
    width = len(grid[0]) if height else 0
    if (not width or not all(len(row) == width for row in grid)
            or not all(math.isfinite(n) for row in grid for n in row)
            or not all(math.isfinite(n) for n in (u, v))):
        raise ValueError("Expected a rectangular finite grid and coordinates")
    if method == "discrete":
        column = min(height - 1, max(0, int(u * width + 0.5)))
        row = min(height - 1, max(0, int(v * height + 0.5)))
        if column >= width:
            raise ValueError("Upstream rectangle clamp yields invalid column")
        return grid[row][column]
    if method != "bilinear":
        raise ValueError("Unknown sampler")
    x, y = u * width - 0.5, v * height - 0.5
    value = 0.0
    for row in (math.floor(y), math.floor(y) + 1):
        for column in (math.floor(x), math.floor(x) + 1):
            if 0 <= row < height and 0 <= column < width:
                weight = (1 - abs(x - column)) * (1 - abs(y - row))
                value += weight * grid[row][column]
    return value
# endregion: sampling


# region: aggregation
def aggregate(values, logits):
    """Pools one head over all levels together using stable softmax."""
    if (not values or len(values) != len(logits)
            or not all(math.isfinite(n) for n in [*values, *logits])):
        raise ValueError("Expected aligned nonempty finite arrays")
    exps = [math.exp(n - max(logits)) for n in logits]
    denominator = sum(exps)
    return sum(value * weight / denominator
               for value, weight in zip(values, exps))
# endregion: aggregation


# region: training
def training_phase(epoch, epochs=72, quiet_epochs=2):
    """Returns the report's augmentation phase for a zero-based epoch.

    This models the paper schedule, not the released R18 YAML overrides.
    """
    if (not all(type(n) is int for n in (epoch, epochs, quiet_epochs))
            or epochs < 1 or not 0 <= epoch < epochs
            or not 0 <= quiet_epochs <= epochs):
        raise ValueError("Epoch must be inside the run")
    return "strong" if epoch < epochs - quiet_epochs else "quiet"
# endregion: training


def main():
    """Checks representative invariants and prints authored calculations."""
    grid = [[1.0, 2.0], [3.0, 4.0]]
    assert sampling_budget([4, 4, 4]) == 86400
    assert sampling_budget([3, 3, 3]) == 64800
    assert box_location([0.5, 0.5, 0.4, 0.2], [4, 0], 4) == [0.7, 0.5]
    assert sample(grid, 0.5, 0.5) == 2.5
    assert sample(grid, 0.5, 0.5, "discrete") == 4.0
    assert sample(grid, 0, 0) == 0.25
    assert aggregate([1, 3], [1000, 1000]) == 2.0
    assert training_phase(69) == "strong"
    assert training_phase(70) == "quiet"
    print(json.dumps({"sample_locations": sampling_budget([4, 4, 4]),
                      "bilinear": sample(grid, 0.5, 0.5),
                      "discrete": sample(grid, 0.5, 0.5, "discrete"),
                      "last_epoch_phase": training_phase(71)}, indent=2))


if __name__ == "__main__":
    main()