# examples/deformable_detr.py
"""Offline scalar Deformable DETR mechanics; authored examples, not inference."""

import argparse
import json
import math


# region: sampling
def sampling_location(reference, offset, shape, points, valid_ratio=(1, 1)):
    """Maps point or box references into a padded level's normalized frame.

    Args:
        reference: Valid-image coordinates [x, y] or [x, y, width, height].
        offset: Authored offset [dx, dy].
        shape: Padded level size [width, height].
        points: Number of samples per level per head.
        valid_ratio: Unpadded fraction [valid_width/W, valid_height/H].

    Returns:
        Two normalized sampling coordinates, with no clamping.
    """
    result = []
    for axis in range(2):
        center = reference[axis] * valid_ratio[axis]
        if len(reference) == 4:
            shift = (offset[axis] / points * reference[axis + 2]
                     * valid_ratio[axis] * 0.5)
        else:
            shift = offset[axis] / shape[axis]
        result.append(center + shift)
    return result
# endregion: sampling


# region: interpolation
def bilinear_sample(grid, u, v):
    """Samples a scalar grid with align_corners=False and zero padding."""
    height, width = len(grid), len(grid[0])
    x, y = u * width - 0.5, v * height - 0.5
    left, top = math.floor(x), math.floor(y)
    total = 0.0
    for row in (top, top + 1):
        for column in (left, left + 1):
            weight = (1 - abs(x - column)) * (1 - abs(y - row))
            if 0 <= row < height and 0 <= column < width:
                total += weight * grid[row][column]
    return total
# endregion: interpolation


# region: aggregation
def normalize_weights(logits):
    """Normalizes one head jointly across levels and points."""
    maximum = max(value for level in logits for value in level)
    exp = [[math.exp(value - maximum) for value in level]
           for level in logits]
    total = sum(sum(level) for level in exp)
    return [[value / total for value in level] for level in exp]


def aggregate(values, weights):
    """Aggregates scalar samples after the value projection, for one head."""
    return sum(value * weights[level][point]
               for level, samples in enumerate(values)
               for point, value in enumerate(samples))
# endregion: aggregation


# region: refinement
def refine_box(reference, delta):
    """Adds reference logits to centers only (2D) or all box axes (4D).

    This arithmetic does not implement the detached autograd boundary.
    """
    result = []
    for axis, change in enumerate(delta):
        if axis < len(reference):
            value = min(1 - 1e-5, max(1e-5, reference[axis]))
            change += math.log(value / (1 - value))
        result.append(1 / (1 + math.exp(-change)))
    return result
# endregion: refinement


# region: proposals
def top_proposals(proposals, count):
    """Ranks valid authored scores without NMS; ties use the slot ID.

    The official two-stage code ranks encoder class channel zero. This
    helper consumes scores already extracted from that ranking channel.
    """
    valid = [proposal for proposal in proposals if proposal['valid']]
    return sorted(valid, key=lambda item: (-item['score'], item['id']))[:count]
# endregion: proposals


def smoke():
    """Returns deterministic mathematical fixtures used for JS parity."""
    weights = normalize_weights([[1000, 1001], [1002, 1003]])
    proposals = [{'id': 0, 'score': .8, 'valid': True},
                 {'id': 1, 'score': .99, 'valid': False},
                 {'id': 2, 'score': .9, 'valid': True}]
    return {
        'samples': [bilinear_sample([[1, 3], [5, 7]], u, v)
                    for u, v in ((.5, .5), (0, 0), (-.5, .5), (1, .25))],
        'point': sampling_location([.5, .5], [2, -1], [8, 4], 4),
        'box': sampling_location([.5, .5, .4, .2], [2, -1], [8, 4], 4),
        'padded': sampling_location([.5, .5], [0, 0], [8, 4], 4, [.75, .5]),
        'weights': weights,
        'aggregate': aggregate([[2, 4], [6, 8]], weights),
        'refined': refine_box([.5, .5, .2, .2], [.1, -.2, .3, -.4]),
        'proposals': [proposal['id'] for proposal in top_proposals(proposals, 2)],
    }


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--smoke', action='store_true',
                        help='Print deterministic arithmetic fixtures as JSON.')
    parser.parse_args()
    print(json.dumps(smoke(), sort_keys=True))