"""Run AE3-TL on a Mac with Core ML, and compare it with the ONNX CPU path.

    python predict_mac.py page.jpg
    python predict_mac.py page.jpg --units all
    python predict_mac.py page.jpg --compare-onnx

The page model runs in Core ML; the relation head stays on the CPU with ONNX Runtime. Table boxes and
caption and footer links are printed. Full cells and text come from the kit's post-processing.
"""
import argparse
import time

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

try:
    import coremltools as ct
except ImportError:
    raise SystemExit("coremltools is required: pip install 'coremltools>=8'")

PACKAGE = "coreml/ae3tl_pages_v3_ep1_ema_fp32.mlpackage"
REL = "weights/ae3tl_relhead_v3_ep1_fp32_opset17.onnx"
PAGES_ONNX = "weights/ae3tl_pages_v3_ep1_fp32_opset17.onnx"
TABLE, CAPTION, FOOTER = (0, 8), (6, 14), (7, 15)
UNITS = {
    "cpu": ct.ComputeUnit.CPU_ONLY,
    "gpu": ct.ComputeUnit.CPU_AND_GPU,
    "all": ct.ComputeUnit.ALL,
    "ane": ct.ComputeUnit.CPU_AND_NE,
}


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


def decode(logits, boxes, W, H, threshold=0.5):
    p = 1.0 / (1.0 + np.exp(-logits[0]))
    order = np.argsort(-p.reshape(-1))[:300]
    scores = p.reshape(-1)[order]
    labels, qids = (order % 16), (order // 16)
    b = boxes[0][qids]
    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 >= threshold
    labels, qids, xyxy = labels[keep], qids[keep], xyxy[keep]
    tables = [(int(q), box.tolist()) for q, box in zip(qids[np.isin(labels, TABLE)], xyxy[np.isin(labels, TABLE)])]
    captions = [int(q) for q in qids[np.isin(labels, CAPTION)]]
    footers = [int(q) for q in qids[np.isin(labels, FOOTER)]]
    return tables, captions, footers


def run_coreml(path, units):
    model = ct.models.MLModel(PACKAGE, compute_units=UNITS[units])
    W, H, x = preprocess(path)
    t0 = time.perf_counter()
    out = model.predict({"images": x})
    dt = time.perf_counter() - t0
    tables, captions, footers = decode(out["logits"], out["boxes"], W, H)
    return tables, captions, footers, out["query_embeddings"], dt


def run_onnx(path):
    session = ort.InferenceSession(PAGES_ONNX, providers=["CPUExecutionProvider"])
    W, H, x = preprocess(path)
    t0 = time.perf_counter()
    out = dict(zip([o.name for o in session.get_outputs()], session.run(None, {"images": x})))
    dt = time.perf_counter() - t0
    tables, captions, footers = decode(out["logits"], out["boxes"], W, H)
    return tables, captions, footers, out["query_embeddings"], dt


def link(embeddings, tables, quarters, rel):
    """The relation head on the CPU: caption or footer query against each table query."""
    result = {}
    for q in quarters:
        pairs = []
        for tq, _ in tables:
            if tq != q:
                pairs += [[int(q), int(tq)], [int(tq), int(q)]]
        if not pairs:
            result[q] = None
            continue
        e = rel.run(None, {"query_embeddings": embeddings[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
        result[q] = (best, round(best_score, 3))
    return result


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("page")
    ap.add_argument("--units", default="gpu", choices=sorted(UNITS))
    ap.add_argument("--compare-onnx", action="store_true")
    args = ap.parse_args()

    rel = ort.InferenceSession(REL, providers=["CPUExecutionProvider"])
    tables, captions, footers, emb, dt = run_coreml(args.page, args.units)
    print(f"compute unit: {args.units}")
    print(f"page model: {dt * 1000:.1f} ms")
    print(f"tables: {len(tables)}")
    for q, box in tables:
        print("  table", q, [round(v) for v in box])
    for kind, quarters in (("caption", captions), ("footer", footers)):
        for q, (best, score) in link(emb, tables, quarters, rel).items():
            print(f"  {kind} {q} -> table {best} score {score}")

    if args.compare_onnx:
        t2, c2, f2, emb2, dt2 = run_onnx(args.page)
        print(f"\nONNX CPU: {dt2 * 1000:.1f} ms, tables {len(t2)}")
        same_tables = len(tables) == len(t2)
        max_dev = 0.0
        for (q1, b1), (q2, b2) in zip(tables, t2):
            max_dev = max(max_dev, max(abs(a - b) for a, b in zip(b1, b2)))
        print(f"same table count: {same_tables}; max box difference: {max_dev:.1f} px")
        l1 = link(emb, tables, captions, rel)
        l2 = link(emb2, t2, c2, rel)
        print(f"same caption links: {l1 == l2}; same footer links: {link(emb, tables, footers, rel) == link(emb2, t2, f2, rel)}")


if __name__ == "__main__":
    main()