import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType


def export_onnx(
    pt_path="models/yolo11-medium-page.pt",
    onnx_path="models/yolo11-medium-page.onnx",
    imgsz=736,
    dynamic=False,
    opset=12,
    half=False,
):
    from ultralytics import YOLO
    print(f"Loading PT model: {pt_path}")
    model = YOLO(pt_path)
    print(f"Exporting to ONNX: {onnx_path}")
    model.export(format="onnx", imgsz=imgsz, dynamic=dynamic, opset=opset, half=half)
    return onnx_path


def downgrade_ir_v10_to_v9(input_path, output_path):
    print(f"Downgrading IR v10 -> v9: {input_path} -> {output_path}")
    model = onnx.load(input_path)
    model.ir_version = 9
    onnx.save(model, output_path)
    return output_path


def quantize_model(input_path, output_path):
    print(f"Quantizing model: {input_path} -> {output_path}")
    quantize_dynamic(
        model_input=input_path,
        model_output=output_path,
        weight_type=QuantType.QUInt8,
    )
    return output_path


def full_pipeline():
    pt_path = "models/yolo11-medium-page-best.pt"
    onnx_path = "models/yolo11-medium-page-best.onnx"
    ir9_path = "models/yolo11-medium-page-best-ir9.onnx"
    quant_path = "models/yolo11-medium-page-best-quant-ir9.onnx"

    export_onnx(pt_path, onnx_path)
    downgrade_ir_v10_to_v9(onnx_path, ir9_path)
    quantize_model(ir9_path, quant_path)

    import os
    size_mb = os.path.getsize(quant_path) / 1024 / 1024
    print(f"Done! Quantized model: {quant_path} ({size_mb:.1f} MB)")


if __name__ == "__main__":
    full_pipeline()