# examples/dino_detr.py
"""Offline DINO detector teaching arithmetic, not a trained detector.

Run with Python 3.9+ to print deterministic JSON. No dependencies or files written.
Sources: https://arxiv.org/html/2203.03605v4 and IDEA-Research/DINO.
"""

import json
import math


# region: queries
def mixed_queries(proposals, content, k):
    """Select encoder boxes but retain learned content by decoder slot.

    Args:
        proposals: Dictionaries with score, box and illustrative feature.
        content: Learned slot vectors, independent of encoder feature ranking.
        k: Number of queries, bounded by both inputs.

    Returns:
        Selected source indices, anchors and learned content vectors.

    Raises:
        ValueError: If k exceeds the available proposals or slots.
    """
    if not isinstance(k, int) or not 0 <= k <= min(len(proposals), len(content)):
        raise ValueError("K must fit proposals and learned content slots")
    ranked = sorted(enumerate(proposals), key=lambda pair: (-pair[1]["score"], pair[0]))
    return [{"source": source, "anchor": list(proposal["box"]),
             "content": list(content[slot])}
            for slot, (source, proposal) in enumerate(ranked[:k])]
# endregion: queries


# region: denoising
def perturb_box(box, scale, negative, magnitude, signs):
    """Perturb corners with half-size-scaled noise and return clipped cxcywh.

    Positive magnitudes use [0, 1]; negatives use [1, 2], before scale.
    Fixed signs/magnitude illustrate the sampler, not its random distribution.
    """
    x, y, width, height = box
    if scale == 0:
        return list(box)
    corners = [x - width / 2, y - height / 2,
               x + width / 2, y + height / 2]
    dimensions = [width / 2, height / 2, width / 2, height / 2]
    noisy = [max(0, min(1, value + sign * size * scale
                        * (magnitude + int(negative))))
             for value, sign, size in zip(corners, signs, dimensions)]
    return [(noisy[0] + noisy[2]) / 2, (noisy[1] + noisy[3]) / 2,
            noisy[2] - noisy[0], noisy[3] - noisy[1]]
# endregion: denoising


# region: mask
def attention_allowed(row, column, groups, group_size):
    """Return permission: row reads column; True means allowed, not masked."""
    dn_count = groups * group_size
    if row >= dn_count:
        return column >= dn_count
    return column >= dn_count or row // group_size == column // group_size
# endregion: mask


# region: refinement
def update_box(reference, delta):
    """Apply a scalar logit-space box update, clipping only logit endpoints."""
    reference = max(1e-6, min(1 - 1e-6, reference))
    return 1 / (1 + math.exp(-(math.log(reference / (1 - reference)) + delta)))
# endregion: refinement


# region: gradients
def gradient_graph(first_offset=0.2, second_weight=1.3, target=0.75):
    """Check an authored adjacent-loss derivative, freezing the reference.

    This is a local finite difference of a graph with stop-gradient semantics.
    Re-evaluating the detached reference after perturbing first_offset would
    accidentally include a path that automatic differentiation blocks.
    """
    first = update_box(0.4, first_offset)
    frozen_reference = first
    delta = second_weight * math.tanh(frozen_reference)
    prediction = update_box(first, delta)
    epsilon = 1e-5

    def loss(value):
        return 0.5 * (value - target) ** 2

    def evaluate(offset):
        return loss(update_box(update_box(0.4, offset), delta))

    gradient = (prediction - target) * prediction * (1 - prediction)
    return {
        "first": first, "prediction": prediction, "loss": loss(prediction),
        "gradient": gradient,
        "finiteDifference": (evaluate(first_offset + epsilon)
                             - evaluate(first_offset - epsilon)) / (2 * epsilon),
        "onceGradient": 0, "referenceGradient": 0,
        "naiveGradient": gradient * (1 + second_weight
                                     * (1 - math.tanh(first) ** 2)
                                     * first * (1 - first)),
    }
# endregion: gradients


def main():
    """Print reproducible fixtures used by the JavaScript parity tests."""
    proposals = [
        {"score": 0.42, "box": [0.2, 0.25, 0.2, 0.25]},
        {"score": 0.91, "box": [0.7, 0.3, 0.25, 0.3]},
        {"score": 0.76, "box": [0.4, 0.7, 0.3, 0.2]},
        {"score": 0.33, "box": [0.75, 0.75, 0.15, 0.2]},
    ]
    box, signs = [0.5, 0.5, 0.2, 0.4], [-1, -1, 1, 1]
    print(json.dumps({
        "queries": mixed_queries(proposals, [[0.2, -0.3], [0.6, 0.1], [-0.1, 0.5]], 2),
        "positive": perturb_box(box, 0.4, False, 0.5, signs),
        "negative": perturb_box(box, 0.4, True, 0.5, signs),
        "mask": [[attention_allowed(r, c, 2, 2) for c in range(6)]
                 for r in range(6)],
        "refinement": update_box(0.25, math.log(3)),
        "gradients": gradient_graph(),
    }, sort_keys=True))


if __name__ == "__main__":
    main()