"""Auto-export best.pt to FP16 ONNX when training finishes or is killed."""
import torch, onnx, os, glob, sys
_orig = torch.load
def _safe(*a, **kw): kw["weights_only"] = False; return _orig(*a, **kw)
torch.load = _safe
from ultralytics import YOLO
from onnxconverter_common import float16

# Find the best model
candidates = glob.glob("runs/detect/train*/weights/best.pt") + \
             glob.glob("train_with_products/run/weights/best.pt")
if not candidates:
    # Try last.pt if best doesn't exist yet
    candidates = glob.glob("runs/detect/train*/weights/last.pt") + \
                 glob.glob("train_with_products/run/weights/last.pt")

if not candidates:
    print("NO WEIGHTS FOUND")
    sys.exit(1)

best = sorted(candidates, key=os.path.getmtime)[-1]
print(f"Exporting: {best} ({os.path.getsize(best)/1024/1024:.0f}MB)")

model = YOLO(best, task="detect")
model.export(format="onnx", opset=17, simplify=True, imgsz=1280)
onnx_path = best.replace(".pt", ".onnx")

m = onnx.load(onnx_path)
m16 = float16.convert_float_to_float16(m, keep_io_types=True)
onnx.save(m16, "/tmp/model_fp16.onnx")
size = os.path.getsize("/tmp/model_fp16.onnx") / 1024 / 1024
print(f"FP16: {size:.1f} MB")
print("EXPORT DONE")