# examples/florence2.py
"""Offline Florence-2 token protocol, not a pretrained vision model.

Authored strict box/text interface; polygon and OCR quadrilateral grammars are
only declared. Quantization and decoder shift conventions follow Microsoft's
processor/model at HF revision 5ca5edf5bd017b9919c05d08aebef5e4c7ac3bac.
No tokenizer, image encoder, checkpoint, model predictions, or network calls.
"""

import argparse
import hashlib
import json
import math
from pathlib import Path
import re

import torch
import torch.nn.functional as functional

REVISION = "5ca5edf5bd017b9919c05d08aebef5e4c7ac3bac"
BINS = 1000
PAD = 1
DECODER_START = 2
LOCATION_OFFSET = 10
VOCABULARY = LOCATION_OFFSET + BINS

# region: task_prompt
# A supported subset of the release processor's task dispatch.
TASKS = {
    "<CAPTION>": ("text", "What does the image describe?"),
    "<OD>": ("boxes", "Locate the objects with category name in the image."),
    "<DENSE_REGION_CAPTION>": (
        "boxes",
        "Locate the objects in the image, with their descriptions.",
    ),
    "<OCR>": ("text", "What is the text in the image?"),
    "<OCR_WITH_REGION>": (
        "ocr_quad",
        "What is the text in the image, with regions?",
    ),
    "<REFERRING_EXPRESSION_SEGMENTATION>": (
        "polygon",
        "Locate {input} in the image with mask",
    ),
}
GRAMMARS = {
    "boxes": "label + four location tokens (x0,y0,x1,y1), repeated",
    "text": "plain text, without spatial tokens",
    "ocr_quad": "text + eight location tokens (four clockwise vertices)",
    "polygon": "clockwise x,y pairs; <sep> parts and <poly>...</poly> groups",
}


def task_prompt(task, text=""):
    """Dispatch a supported task and construct its released text prompt.

    Args:
        task: Exact task marker from TASKS.
        text: Referring expression for the segmentation task only.

    Returns:
        A dictionary with task, prompt, format, and declared grammar.

    Raises:
        ValueError: For unsupported tasks or incorrect task input.
    """
    if not isinstance(task, str) or task not in TASKS:
        raise ValueError("Unsupported task marker.")
    if not isinstance(text, str):
        raise ValueError("Task input must be text.")
    kind, template = TASKS[task]
    needs_input = "{input}" in template
    if needs_input and (not text.strip() or "<" in text or ">" in text):
        raise ValueError("A nonempty referring expression is required.")
    if not needs_input and text:
        raise ValueError("This task takes no additional text.")
    return {
        "task": task,
        "prompt": template.format(input=text),
        "format": kind,
        "grammar": GRAMMARS[kind],
    }


# endregion: task_prompt


def image_scale(size):
    """Validate (width, height) and return x,y,x,y axis scales on CPU."""
    if not isinstance(size, (tuple, list)) or len(size) != 2:
        raise ValueError("Image size must be (width, height).")
    if any(
        isinstance(value, bool)
        or not isinstance(value, (int, float))
        or not math.isfinite(value)
        or value <= 0
        for value in size
    ):
        raise ValueError("Image dimensions must be finite and positive.")
    return torch.tensor([*size, *size], dtype=torch.float64)


def box_tensor(boxes):
    """Validate finite CPU [N,4] pixel boxes in corner order."""
    values = torch.as_tensor(boxes)
    if values.device.type != "cpu" or values.dtype == torch.bool:
        raise ValueError("Boxes must be numeric CPU values.")
    if values.is_complex() or values.ndim != 2 or values.shape[1] != 4:
        raise ValueError("Boxes must have shape [N,4].")
    values = values.to(torch.float64)
    if not torch.isfinite(values).all():
        raise ValueError("Boxes must be finite.")
    if (values[:, 2:] < values[:, :2]).any():
        raise ValueError("Box corners must satisfy x0<=x1 and y0<=y1.")
    return values


# region: quantize
def quantize(boxes, size):
    """Floor and clamp [N,4] boxes using width for x and height for y.

    Coordinates outside the image are clamped, matching BoxQuantizer. The
    half-bin error bound applies to coordinates inside the closed image.
    """
    scales = image_scale(size)
    values = box_tensor(boxes)
    return (values / (scales / BINS)).floor().clamp(0, BINS - 1).long()


# endregion: quantize


# region: serialize
def serialize(labels, boxes, size, task="<OD>"):
    """Serialize annotated boxes, not inferred boxes, as label + four tokens.

    This strict subset requires nonempty labels without token delimiters.
    Polygon parts and OCR quadrilaterals have separate grammars and are not
    accepted by this four-coordinate serializer. No BPE tokenizer is modeled.
    """
    if task_prompt(task)["format"] != "boxes":
        raise ValueError("The box serializer requires a box task.")
    bins = quantize(boxes, size)
    if not isinstance(labels, (list, tuple)) or len(labels) != len(bins):
        raise ValueError("One label is required per box.")
    pieces = []
    for label, corners in zip(labels, bins.tolist()):
        if (
            not isinstance(label, str)
            or not label.strip()
            or "<" in label
            or ">" in label
            or label != label.strip()
        ):
            raise ValueError("Labels must be stripped text without brackets.")
        pieces.append(label + "".join(f"<loc_{bin_id}>" for bin_id in corners))
    return "".join(pieces)


# endregion: serialize


# region: teacher_forcing
def teacher_forcing(target_ids, logits):
    """Right-shift decoder inputs once; score same-position target labels.

    Args:
        target_ids: CPU int64 [B,T] IDs, with PAD=1 only in suffix padding.
        logits: Finite CPU floating [B,T,V] fixed prediction logits.

    Returns:
        Decoder input IDs, labels with PAD replaced by -100, and mean CE.
        All-padding batches return differentiable zero (authored extension).

    Raises:
        ValueError: For invalid shapes, types, IDs, or padding placement.
    """
    if (
        not isinstance(target_ids, torch.Tensor)
        or target_ids.dtype != torch.long
        or target_ids.ndim != 2
        or min(target_ids.shape) < 1
        or target_ids.device.type != "cpu"
    ):
        raise ValueError("Targets must be nonempty CPU int64 [B,T].")
    if (
        not isinstance(logits, torch.Tensor)
        or logits.ndim != 3
        or logits.shape[:2] != target_ids.shape
        or logits.shape[2] <= 2
        or not logits.is_floating_point()
        or logits.device.type != "cpu"
        or not torch.isfinite(logits).all()
    ):
        raise ValueError("Logits must be finite CPU float [B,T,V].")
    if (target_ids < 0).any() or (target_ids >= logits.shape[2]).any():
        raise ValueError("Target IDs must belong to the logit vocabulary.")
    padding = target_ids == PAD
    if (padding[:, :-1] & ~padding[:, 1:]).any():
        raise ValueError("Padding must be a suffix.")
    labels = target_ids.masked_fill(padding, -100)
    decoder_inputs = torch.full_like(target_ids, PAD)
    decoder_inputs[:, 0] = DECODER_START
    decoder_inputs[:, 1:] = labels[:, :-1]
    decoder_inputs.masked_fill_(decoder_inputs == -100, PAD)
    if (~padding).any():
        loss = functional.cross_entropy(
            logits.reshape(-1, logits.shape[-1]),
            labels.reshape(-1),
            ignore_index=-100,
            reduction="mean",
        )
    else:
        loss = logits.sum() * 0
    return decoder_inputs, labels, loss


# endregion: teacher_forcing


# region: parse_dequantize
def dequantize(bins, size):
    """Decode integer [N,4] bins to centers in original image coordinates."""
    scales = image_scale(size)
    values = torch.as_tensor(bins)
    if (
        values.device.type != "cpu"
        or values.ndim != 2
        or values.shape[1] != 4
        or values.dtype == torch.bool
        or values.is_floating_point()
        or values.is_complex()
        or (values < 0).any()
        or (values >= BINS).any()
    ):
        raise ValueError("Location bins must be integer [N,4] in [0,999].")
    decoded = (values.to(torch.float64) + 0.5) * (scales / BINS)
    return box_tensor(decoded)


def parse_dequantize(sequence, task, size):
    """Parse the complete box/text body; reject every malformed remainder.

    Unlike the release regex parser, this authored strict interface never
    skips unmatched spans or accepts out-of-range bins. Caller supplies the
    body only, without BOS/EOS/PAD wrappers. Polygon/OCR-region formats are
    declared but deliberately unsupported, rather than read as boxes.
    """
    kind = task_prompt(
        task, "object" if task == "<REFERRING_EXPRESSION_SEGMENTATION>" else ""
    )["format"]
    image_scale(size)
    if not isinstance(sequence, str):
        raise ValueError("Sequence body must be text.")
    if kind in ("polygon", "ocr_quad"):
        raise NotImplementedError("This interface only parses boxes/text.")
    if kind == "text":
        if "<" in sequence or ">" in sequence:
            raise ValueError("Text output cannot contain special tokens.")
        return {task: sequence}
    record = re.compile(r"([^<>]+)((?:<loc_[0-9]+>){4})")
    labels, bins, cursor = [], [], 0
    while cursor < len(sequence):
        match = record.match(sequence, cursor)
        if match is None or match[1] != match[1].strip():
            raise ValueError("Malformed label/box sequence.")
        labels.append(match[1])
        bins.append([int(value) for value in re.findall(r"[0-9]+", match[2])])
        cursor = match.end()
    corners = torch.tensor(bins, dtype=torch.long).reshape(-1, 4)
    boxes = dequantize(corners, size).tolist()
    return {task: {"bboxes": boxes, "labels": labels}}


# endregion: parse_dequantize


def smoke(seed=7):
    """Execute seeded fixed-logit training algebra and annotated fixtures."""
    size = (640, 360)
    pixels = torch.tensor(
        [[0.0, 0.0, 640.0, 360.0], [64.0, 72.0, 480.0, 270.0]],
        dtype=torch.float64,
    )
    bins = quantize(pixels, size)
    sequence = serialize(["person", "cart"], pixels, size)
    parsed = parse_dequantize(sequence, "<OD>", size)
    target_ids = torch.tensor(
        [
            [
                3,
                *(LOCATION_OFFSET + bins[0]).tolist(),
                4,
                *(LOCATION_OFFSET + bins[1]).tolist(),
                2,
            ],
            [5, 6, 2, *([PAD] * 8)],
        ]
    )
    generator = torch.Generator().manual_seed(seed)
    logits = torch.randn(
        2,
        11,
        VOCABULARY,
        generator=generator,
        dtype=torch.float64,
        requires_grad=True,
    )
    decoder_inputs, labels, loss = teacher_forcing(target_ids, logits)
    loss.backward()
    errors = (dequantize(bins, size) - pixels).abs()
    return {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "revision": REVISION,
        "seed": seed,
        "boundary": "offline CPU protocol; annotated boxes and fixed logits",
        "reductions": {
            "bins_per_axis": BINS,
            "batch": 2,
            "vocabulary": VOCABULARY,
            "sequence_length": 11,
            "loss": "mean over non-PAD targets",
        },
        "shapes": {
            "boxes": list(pixels.shape),
            "targets": list(target_ids.shape),
            "logits": list(logits.shape),
        },
        "image_size": list(size),
        "dispatch": task_prompt("<OD>"),
        "serialized": sequence,
        "parsed": parsed,
        "quantized": bins.tolist(),
        "decoder_inputs": decoder_inputs.tolist(),
        "labels": labels.tolist(),
        "valid_targets": int((labels != -100).sum()),
        "loss": float(loss.detach()),
        "gradient_norm": float(logits.grad.norm()),
        "ignored_gradient_max": float(logits.grad[labels == -100].abs().max()),
        "max_error_xy": errors.amax(dim=0).tolist(),
        "half_bin_xy": (image_scale(size) / (2 * BINS)).tolist(),
    }


def main():
    """Print finite JSON only when --smoke is explicitly requested."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--seed", type=int, default=7)
    arguments = parser.parse_args()
    if not arguments.smoke:
        parser.error("Use --smoke for the offline CPU protocol reference.")
    print(json.dumps(smoke(arguments.seed), indent=2, allow_nan=False))


if __name__ == "__main__":
    main()