"""Quantization settings for the published decoder; this script never retrains it."""

import argparse
import hashlib
from pathlib import Path

from onnxruntime.quantization import QuantType, quantize_dynamic


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path, help="Original FP32 decoder_model.onnx")
    parser.add_argument("output", type=Path, help="New INT8 decoder output path")
    args = parser.parse_args()
    expected = "6b1fb216d542c4b2a4fa5b9d7ae3522081eb85fb959d2cecd28055af956a8a5e"
    actual = hashlib.sha256(args.input.read_bytes()).hexdigest()
    if actual != expected:
        parser.error(f"Input SHA-256 mismatch: {actual}")
    if args.output.exists():
        parser.error("Output already exists; choose a new path")
    quantize_dynamic(
        str(args.input), str(args.output),
        op_types_to_quantize=["MatMul"],
        weight_type=QuantType.QInt8,
        per_channel=True,
    )
    print(hashlib.sha256(args.output.read_bytes()).hexdigest(), args.output.name)


if __name__ == "__main__":
    main()