#!/usr/bin/env python3
"""TableFormerV2 ONNX inference — table structure recognition.

Self-contained: requires only numpy, onnxruntime and Pillow.

    python inference.py <table_image.png>

Pipeline: encoder.onnx -> greedy loop over decoder.onnx -> bbox_head.onnx.
The decoder graph is cache-free and re-runs the whole prefix each step
(vocab is 13 tokens, sequences are short, so this is cheap).

Outputs OTSL structure tokens and one bbox per data cell, xyxy normalized
to [0, 1] relative to the input table crop.
"""

import sys
from pathlib import Path

import numpy as np
import onnxruntime as ort
from PIL import Image

HERE = Path(__file__).parent

ID2TOKEN = [
    "<pad>", "[UNK]", "<start>", "<end>", "<ecel>", "<fcel>", "<lcel>",
    "<ucel>", "<xcel>", "<nl>", "<ched>", "<rhed>", "<srow>",
]
BOS_ID, EOS_ID = 2, 3
DATA_CELL_IDS = {4, 5, 10, 11, 12}  # ecel, fcel, ched, rhed, srow

IMAGE_SIZE = 448
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)  # ImageNet
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)


def preprocess(image: Image.Image) -> np.ndarray:
    """RGB table crop -> (1, 3, 448, 448) float32, ImageNet-normalized."""
    image = image.convert("RGB").resize((IMAGE_SIZE, IMAGE_SIZE), Image.BILINEAR)
    x = np.asarray(image, dtype=np.float32) / 255.0
    x = (x - MEAN) / STD
    return x.transpose(2, 0, 1)[None]


class TableFormerV2Onnx:
    def __init__(self, model_dir=HERE, providers=("CPUExecutionProvider",)):
        model_dir = Path(model_dir)
        p = list(providers)
        self.encoder = ort.InferenceSession(str(model_dir / "encoder.onnx"), providers=p)
        self.decoder = ort.InferenceSession(str(model_dir / "decoder.onnx"), providers=p)
        self.bbox_head = ort.InferenceSession(str(model_dir / "bbox_head.onnx"), providers=p)

    def predict(self, image: Image.Image, max_length: int = 512):
        """Returns (otsl_tokens, bboxes). bboxes: (num_cells, 4) xyxy in [0, 1]."""
        images = preprocess(image)
        (enc_hidden,) = self.encoder.run(None, {"images": images})

        # Greedy generation
        ids = np.array([[BOS_ID]], dtype=np.int64)
        for _ in range(max_length):
            logits, _ = self.decoder.run(
                None, {"input_ids": ids, "encoder_hidden": enc_hidden}
            )
            next_id = int(logits[0, -1].argmax())
            ids = np.concatenate([ids, [[next_id]]], axis=1)
            if next_id == EOS_ID:
                break

        # Hidden states of the full sequence -> bboxes at data-cell positions
        _, hidden = self.decoder.run(
            None, {"input_ids": ids, "encoder_hidden": enc_hidden}
        )
        cell_pos = [i for i, t in enumerate(ids[0].tolist()) if t in DATA_CELL_IDS]
        if cell_pos:
            (bboxes,) = self.bbox_head.run(
                None,
                {"cell_embeddings": hidden[0, cell_pos], "encoder_hidden": enc_hidden},
            )
        else:
            bboxes = np.zeros((0, 4), dtype=np.float32)

        tokens = [
            ID2TOKEN[t] for t in ids[0].tolist() if t not in (BOS_ID, EOS_ID, 0)
        ]
        return tokens, bboxes


if __name__ == "__main__":
    if len(sys.argv) != 2:
        sys.exit(f"usage: {sys.argv[0]} <table_image>")
    img = Image.open(sys.argv[1])
    model = TableFormerV2Onnx()
    tokens, bboxes = model.predict(img)

    print("OTSL:", " ".join(tokens))
    print(f"{len(bboxes)} cells (xyxy in original image pixels):")
    scale = np.array([img.width, img.height, img.width, img.height])
    for b in bboxes * scale:
        print(f"  ({b[0]:7.1f}, {b[1]:7.1f}, {b[2]:7.1f}, {b[3]:7.1f})")