"""Minimal AE3-TL example: page model and relation head with ONNX Runtime.

Prints the table boxes and, for each caption and footer box, the table it links to.
For full tables with cells and text, run extract.py, which uses src/ae3_tl.
"""
import os
import sys
import numpy as np
from PIL import Image
import onnxruntime as ort

PAGES = "weights/ae3tl_pages_v3_ep1_fp32_opset17.onnx"
REL = "weights/ae3tl_relhead_v3_ep1_fp32_opset17.onnx"
TABLE, CAPTION, FOOTER = (0, 8), (6, 14), (7, 15)


def _provider_list():
    """Pick the fastest provider this machine has. AE3_PROVIDER overrides the choice."""
    forced = os.environ.get("AE3_PROVIDER", "auto")
    have = ort.get_available_providers()
    if forced in ("cpu", "cuda", "coreml"):
        choice = forced
    elif "CUDAExecutionProvider" in have:
        choice = "cuda"
    elif "CoreMLExecutionProvider" in have:
        choice = "coreml"
    else:
        choice = "cpu"
    print(f"provider: {choice} (available: {', '.join(have)})")
    if choice == "coreml":
        cache = os.environ.get("AE3_COREML_CACHE", os.path.expanduser("~/.cache/ae3tl-coreml"))
        return [("CoreMLExecutionProvider", {"ModelFormat": "MLProgram", "MLComputeUnits": "ALL",
                 "RequireStaticInputShapes": "1", "ModelCacheDirectory": cache}),
                "CPUExecutionProvider"]
    if choice == "cuda":
        return ["CUDAExecutionProvider", "CPUExecutionProvider"]
    return ["CPUExecutionProvider"]


def _session(path, prov):
    """Create a session. If CoreML is chosen it must start, or we stop with the reason."""
    coreml = any("CoreML" in str(p) for p in prov)
    model, opts = path, None
    if coreml:
        if "CoreMLExecutionProvider" not in ort.get_available_providers():
            sys.exit("CoreMLExecutionProvider is missing from this onnxruntime build. On Apple "
                     "Silicon install the standard wheel: pip install -U 'onnxruntime>=1.20'.")
        b1 = os.path.splitext(path)[0] + "_b1.onnx"
        if os.path.exists(b1):
            model = b1
        else:
            opts = ort.SessionOptions()
            opts.add_free_dimension_override_by_name("batch", 1)
    try:
        return (ort.InferenceSession(model, opts, providers=prov) if opts
                else ort.InferenceSession(model, providers=prov))
    except Exception as exc:
        if not coreml:
            raise
        sys.exit(f"CoreML failed to start: {type(exc).__name__}: {exc}")


def main(path):
    prov = _provider_list()
    pages = _session(PAGES, prov)
    rel = _session(REL, prov)

    img = Image.open(path).convert("RGB")
    W, H = img.size
    x = np.asarray(img.resize((768, 768), Image.BILINEAR), dtype=np.float32) / 255.0
    x = x.transpose(2, 0, 1)[None]

    logits, boxes, emb, _ = pages.run(None, {"images": x})
    p = 1.0 / (1.0 + np.exp(-logits[0]))                 # queries x 16 classes
    order = np.argsort(-p.reshape(-1))[:300]
    scores = p.reshape(-1)[order]
    labels, qids = (order % 16), (order // 16)
    b = boxes[0][qids]                                   # cx, cy, w, h, normalised
    xyxy = np.stack([b[:, 0] - b[:, 2] / 2, b[:, 1] - b[:, 3] / 2,
                     b[:, 0] + b[:, 2] / 2, b[:, 1] + b[:, 3] / 2], 1)
    xyxy = xyxy * np.array([W, H, W, H], dtype=np.float32)
    keep = scores >= 0.5
    labels, qids, xyxy = labels[keep], qids[keep], xyxy[keep]

    tab = np.isin(labels, TABLE)
    tables = [(int(q), box.tolist()) for q, box in zip(qids[tab], xyxy[tab])]
    print("tables:", len(tables))
    for q, box in tables:
        print("  table", q, [round(v) for v in box])

    for kind, classes in (("caption", CAPTION), ("footer", FOOTER)):
        sel = np.isin(labels, classes)
        for q in qids[sel]:
            pairs = []
            for tq, _ in tables:
                if tq != q:
                    pairs += [[int(q), int(tq)], [int(tq), int(q)]]
            if not pairs:
                print(kind, q, "no tables"); continue
            e = rel.run(None, {"query_embeddings": emb[0],
                               "pairs": np.asarray(pairs, dtype=np.int64)})[0]
            e = 1.0 / (1.0 + np.exp(-e))
            best, best_score = None, -1.0
            for k, (tq, _) in enumerate(tables):
                m = float((e[2 * k] + e[2 * k + 1]) / 2.0)
                if m > best_score:
                    best, best_score = tq, m
            print(kind, q, "-> table", best, "score", round(best_score, 3))


if __name__ == "__main__":
    main(sys.argv[1])