#!/usr/bin/env python3
"""Convert SCRFD det_10g.onnx from FP32 to FP16.

Usage:
    pip install onnx onnxconverter-common
    python scripts/convert_scrfd_fp16.py \
        --input /path/to/det_10g.onnx \
        --output /path/to/det_10g_fp16.onnx

Design decisions:
- op_block_list=['BatchNormalization'] — epsilon 1e-5 underflows to 0 in FP16 → NaN.
  Keeping BN in FP32 prevents this while still converting ~95% of ops to FP16.
- keep_io_types=True — Input/output remain float32 for compatibility.
  No preprocessing changes needed in SCRFD pipeline.
- onnx.checker validates structural integrity after conversion.
"""

import argparse
import sys
from pathlib import Path


def convert_fp16(input_path: str, output_path: str) -> None:
    """Convert ONNX model from FP32 to FP16."""
    try:
        import onnx
        from onnxconverter_common import float16
    except ImportError:
        print("Missing dependencies. Install:")
        print("  pip install onnx onnxconverter-common")
        sys.exit(1)

    input_file = Path(input_path)
    if not input_file.exists():
        print(f"Input file not found: {input_path}")
        sys.exit(1)

    print(f"Loading {input_path} ...")
    model = onnx.load(input_path)

    input_size_mb = input_file.stat().st_size / (1024 * 1024)
    print(f"  Input size: {input_size_mb:.1f} MB")
    print(f"  Opset version: {model.opset_import[0].version}")

    # Convert to FP16 with BatchNormalization excluded
    # BN epsilon (1e-5) underflows to 0 in FP16 → division by zero → NaN
    print("Converting to FP16 (excluding BatchNormalization) ...")
    model_fp16 = float16.convert_float_to_float16(
        model,
        op_block_list=["BatchNormalization"],
        keep_io_types=True,
    )

    # Validate
    print("Validating converted model ...")
    onnx.checker.check_model(model_fp16)

    # Save
    output_file = Path(output_path)
    output_file.parent.mkdir(parents=True, exist_ok=True)
    onnx.save(model_fp16, output_path)

    output_size_mb = output_file.stat().st_size / (1024 * 1024)
    ratio = output_size_mb / input_size_mb * 100
    print(f"  Output size: {output_size_mb:.1f} MB ({ratio:.0f}% of original)")
    print(f"  Saved to: {output_path}")
    print("Done.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert SCRFD det_10g.onnx FP32 → FP16")
    parser.add_argument("--input", required=True, help="Path to FP32 det_10g.onnx")
    parser.add_argument("--output", required=True, help="Output path for FP16 model")
    args = parser.parse_args()

    convert_fp16(args.input, args.output)