# examples/lw_detr.py
"""Calculate LW-DETR teaching examples using only the standard library.

These helpers implement indexing and forward arithmetic, not a trained detector.
Run `python3 examples/lw_detr.py` to print the deterministic reference fixture.
"""

import json
import math


# region: windows
def window_order(rows, cols, side):
    """Return row-major token IDs in window-major storage order."""
    if any(not isinstance(v, int) or isinstance(v, bool) or v <= 0
           for v in (rows, cols, side)):
        raise ValueError("Expected positive integer dimensions")
    if rows % side or cols % side:
        raise ValueError("Grid must be divisible by window side")
    return [(y + dy) * cols + x + dx
            for y in range(0, rows, side)
            for x in range(0, cols, side)
            for dy in range(side) for dx in range(side)]


def restore_order(values, order):
    """Undo a bijective storage permutation without changing token values."""
    if len(values) != len(order) or sorted(order) != list(range(len(order))):
        raise ValueError("Expected a permutation")
    restored = [None] * len(values)
    for index, token_id in enumerate(order):
        restored[token_id] = values[index]
    return restored
# endregion: windows


# region: connectivity
def attention_pairs(rows, cols, side):
    """Count ordered query/key pairs per head, including self-attention."""
    window_order(rows, cols, side)
    count = rows * cols
    return {"local": count * side * side, "global": count * count}


def reachable_tokens(rows, cols, side, selected, layers):
    """Compute structural reach, not learned attention weights."""
    window_order(rows, cols, side)
    if not isinstance(selected, int) or not 0 <= selected < rows * cols:
        raise ValueError("Invalid token")
    reached = {selected}
    for layer in layers:
        if layer == "global":
            reached = set(range(rows * cols))
        elif layer == "local":
            reached = {
                ((token // cols) // side * side + dy) * cols
                + (token % cols) // side * side + dx
                for token in reached for dy in range(side) for dx in range(side)
            }
        else:
            raise ValueError("Unknown attention kind")
    return sorted(reached)
# endregion: connectivity


# region: projector
def projector_geometry(variant, resolution):
    """Describe output geometry; no C2f or feature fusion is executed."""
    if variant not in ("tiny", "small", "medium", "large", "xlarge"):
        raise ValueError("Unknown variant")
    large = variant in ("large", "xlarge")
    divisor = 32 if large else 16
    if not isinstance(resolution, int) or resolution <= 0 or resolution % divisor:
        raise ValueError("Resolution must divide exactly")
    return [{"stride": stride, "side": resolution // stride,
             "channels": 384 if large else 256}
            for stride in ((8, 32) if large else (16,))]
# endregion: projector


# region: queries
def query_budget(variant, training):
    """Count Group DETR query slots sharing decoder weights."""
    if variant not in ("tiny", "small", "medium", "large", "xlarge"):
        raise ValueError("Unknown variant")
    per_group = 100 if variant == "tiny" else 300
    groups = 13 if training else 1
    return {"groups": groups, "perGroup": per_group,
            "total": groups * per_group, "layers": 3}
# endregion: queries


# region: loss
def iou_target(score, iou):
    """Evaluate the paper's positive target t=s^.25*u^.75."""
    if any(not math.isfinite(v) or not 0 <= v <= 1 for v in (score, iou)):
        raise ValueError("Expected probability in [0,1]")
    return score ** 0.25 * iou ** 0.75


def ia_bce(score, iou, positive):
    """Return paper-equation forward loss, without autograd or matching.

    Endpoints use extended-real BCE with 0*log(0)=0. The official criterion
    additionally controls detached targets, normalization and reduction.
    """
    target = iou_target(score, iou) if positive else 0.0
    if not math.isfinite(score) or not 0 <= score <= 1:
        raise ValueError("Expected probability in [0,1]")
    positive_term = (0.0 if target == 0 else
                     math.inf if score == 0 else -target * math.log(score))
    negative_term = (0.0 if target == 1 else
                     math.inf if score == 1 else
                     -(1 - target) * math.log1p(-score))
    return (1 if positive else score * score) * (positive_term + negative_term)
# endregion: loss


def main():
    """Print the deterministic fixture compared against JavaScript in tests."""
    order = window_order(4, 4, 2)
    print(json.dumps({
        "windows": order,
        "restored": restore_order(order, order),
        "connectivity": attention_pairs(4, 4, 2),
        "reach": reachable_tokens(4, 4, 2, 0, ["local", "global", "local"]),
        "projector": projector_geometry("large", 640),
        "queries": query_budget("tiny", True),
        "target": iou_target(0.16, 0.81),
        "negative": ia_bce(0.5, 0, False),
    }, sort_keys=True))


if __name__ == "__main__":
    main()