"""Create an optional INT8 ONNX variant; evaluate it before selecting for serving."""

import argparse
import json
import shutil
from pathlib import Path
from onnxruntime.quantization import quantize_dynamic, QuantType
from image_to_latex.artifacts import digest


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--source", type=Path, default=Path("artifacts/latex-ocr"))
    p.add_argument("--output", type=Path, default=Path("artifacts/latex-ocr-int8"))
    args = p.parse_args()
    if args.output.exists():
        p.error("Output must be new; original artifacts will not be overwritten")
    args.output.mkdir(parents=True)
    for name in ["encoder", "decoder"]:
        quantize_dynamic(
            str(args.source / f"{name}.onnx"),
            str(args.output / f"{name}.int8.onnx"),
            weight_type=QuantType.QInt8,
            per_channel=True,
            op_types_to_quantize=["MatMul", "Gemm"],
            extra_options={"MatMulConstBOnly": True},
        )
    for name in ["image_resizer.onnx", "tokenizer.json"]:
        shutil.copy2(args.source / name, args.output / name)
    metadata = {
        p.name: dict(bytes=p.stat().st_size, sha256=digest(p))
        for p in args.output.iterdir()
        if p.is_file()
    }
    (args.output / "quantization.json").write_text(
        json.dumps(metadata, indent=2) + "\n"
    )
    print(json.dumps(metadata, indent=2))


if __name__ == "__main__":
    main()