#!/usr/bin/env python3
# Copyright (c) 2026, CIIL Anuvani-LID. Apache-2.0.
"""How much of each accelerator encoder the phone accelerators can take, by ONNX Runtime's own mobile-usability
check (python -m onnxruntime.tools.check_onnx_model_mobile_usability): for NNAPI, Core ML NeuralNetwork and Core ML
MLProgram, the number of partitions, the nodes they cover and the operators they cannot run. One partition covering
every node means the accelerator runs the whole encoder in one call; several mean round trips through the CPU.

The check reads the graph only. It cannot see a phone's driver, and it lists caveats it does not verify (e.g. NNAPI
takes only 2-D Conv: a graph with 1-D convolutions passes the check and still splits on a device). What a runtime
actually did is in results/android/emulator_x86_64.json (NNAPI, "placement") and results/gates/qnn_soc*.json (QNN).

    python3 gate/coverage.py [models/anuvani_lid.npu-15s.encoder.onnx ...]     (default: every *.encoder.onnx)
Report -> results/gates/accelerator_coverage.json; the checker's full output -> results/gates/usability/<model>.txt
"""
import glob
import json
import os
import re
import subprocess
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "tools"))
import common  # noqa: E402

SECTIONS = {"NNAPI": "nnapi", "CoreML NeuralNetwork": "coreml_neuralnetwork", "CoreML MLProgram": "coreml_mlprogram"}


def parse(text):
    out = {}
    for title, key in SECTIONS.items():
        m = re.search(rf"Checking {title}\n(.*?)(?:={{8,}}|\Z)", text, re.S)
        if not m:
            continue
        sec = m.group(1)
        p = re.search(r"(\d+) partitions with a total of (\d+)/(\d+) nodes", sec)
        u = re.search(r"Unsupported ops: (.*)", sec)
        ok = re.search(r"Model should perform well with .* as is: (YES|NO)", sec)
        out[key] = {"partitions": int(p.group(1)) if p else None,
                    "nodes_supported": int(p.group(2)) if p else None,
                    "nodes_total": int(p.group(3)) if p else None,
                    "unsupported_ops": [o.strip() for o in u.group(1).split(",")] if u else [],
                    "checker_verdict_perform_well": ok.group(1) if ok else None}
    return out


def main():
    models = sys.argv[1:] or sorted(glob.glob(os.path.join(common.REPO, "models", "*.encoder.onnx")))
    import onnxruntime as ort
    rep = {"_doc": __doc__.strip().split("\n\n")[0], "onnxruntime": ort.__version__, "models": {}}
    os.makedirs(os.path.join(common.REPO, "results", "gates", "usability"), exist_ok=True)
    for m in models:
        r = subprocess.run([sys.executable, "-m", "onnxruntime.tools.check_onnx_model_mobile_usability", m],
                           capture_output=True, text=True)
        text = r.stdout + r.stderr
        text = re.sub(r"^INFO:\s+", "", text, flags=re.M)
        name = os.path.basename(m)
        with open(os.path.join(common.REPO, "results", "gates", "usability", name[:-5] + ".txt"), "w") as f:
            f.write(text)
        rep["models"][name] = {"sha256": common.sha256(m), **parse(text)}
        print(name, json.dumps(rep["models"][name]))
    out = os.path.join(common.REPO, "results", "gates", "accelerator_coverage.json")
    json.dump(rep, open(out, "w"), indent=1)
    print("->", os.path.relpath(out, common.REPO))


if __name__ == "__main__":
    main()