#!/usr/bin/env python3
"""Export the pinned #000049 shadow-NLI checkpoint to ONNX (+ optional
int8 dynamic quantization) — the §3 CPU speedup path.

`onnxruntime` on a quantized cross-encoder is typically 2–4× faster than
torch on CPU and drops the torch forward path; `ShadowNLI._ensure_loaded`
will prefer the export automatically if it finds it at the conventional
location (`~/.arborist/models/nli/<model_version>/onnx/`, or wherever
`ARBORIST_NLI_ONNX_DIR` points). Run once after `make bootstrap-nli`:

    python3 bench/scripts/export_nli_onnx.py            # export + int8 quantize
    python3 bench/scripts/export_nli_onnx.py --no-quantize
    python3 bench/scripts/export_nli_onnx.py --out /some/dir

Requires `optimum[onnxruntime]` (in the `[nli]` extra). SHADOW
infrastructure — the ONNX model is the same pinned checkpoint, same
labels; nothing about audit_mode changes (cf. ticket #000049 §7).
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parents[2]


def main(argv=None) -> int:
    sys.path.insert(0, str(REPO))
    from arborist.qa.nli.shadow import load_manifest, ShadowNLI

    manifest = load_manifest()
    repo_id = manifest["hf_repo"]
    rev = manifest.get("pinned_revision")
    default_out = ShadowNLI(manifest)._onnx_dir()

    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, default=default_out,
                    help=f"output dir (default: {default_out})")
    ap.add_argument("--no-quantize", action="store_true", help="skip int8 dynamic quantization")
    args = ap.parse_args(argv)

    try:
        from transformers import AutoTokenizer
        from optimum.onnxruntime import ORTModelForSequenceClassification
        if not args.no_quantize:
            from optimum.onnxruntime import ORTQuantizer
            from optimum.onnxruntime.configuration import AutoQuantizationConfig
    except ImportError as e:
        print(f"[export-nli-onnx] missing dependency: {e}\n"
              f"  install with: pip install 'arborist[nli]'  (pulls optimum[onnxruntime])",
              file=sys.stderr)
        return 2

    out: Path = args.out
    out.mkdir(parents=True, exist_ok=True)
    print(f"[export-nli-onnx] {repo_id}@{rev or 'main'} → {out}  (quantize={'no' if args.no_quantize else 'int8-dynamic'})", flush=True)

    tok = AutoTokenizer.from_pretrained(repo_id, revision=rev)
    model = ORTModelForSequenceClassification.from_pretrained(repo_id, revision=rev, export=True)
    model.save_pretrained(out)
    tok.save_pretrained(out)
    print(f"[export-nli-onnx] fp32 ONNX written ({sum(f.stat().st_size for f in out.glob('*.onnx'))/1e6:.1f} MB of .onnx)", flush=True)

    if not args.no_quantize:
        quantizer = ORTQuantizer.from_pretrained(out)
        qconfig = AutoQuantizationConfig.avx2(is_static=False, per_channel=False)
        quantizer.quantize(save_dir=out, quantization_config=qconfig)
        print(f"[export-nli-onnx] int8 quantized; dir now {sum(f.stat().st_size for f in out.rglob('*.onnx'))/1e6:.1f} MB of .onnx", flush=True)

    # sanity: load via ShadowNLI's ONNX path and run one pair
    import os
    os.environ["ARBORIST_NLI_ONNX_DIR"] = str(out)
    n = ShadowNLI(manifest)
    n._ensure_loaded()
    if not n.available or not (n.backend or "").startswith("onnx"):
        print(f"[export-nli-onnx] WARNING: ShadowNLI did not pick up the ONNX export "
              f"(available={n.available}, backend={n.backend}, reason={n._reason})", file=sys.stderr)
        return 1
    pe, pn, pc = n._nli("Jupiter is the largest planet.", "Mercury is the largest planet.")
    print(f"[export-nli-onnx] sanity OK — backend={n.backend} device={n.device}; "
          f"NLI(Jupiter-is-largest, Mercury-is-largest) → contradiction={pc:.3f} entail={pe:.3f}")
    print(f"[export-nli-onnx] done. ShadowNLI will now auto-prefer {out} (or set ARBORIST_NLI_ONNX_DIR).")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
