#!/usr/bin/env python3
"""
macOS에서 실행: ONNX → CoreML FP16 변환

사용법:
    python convert_coreml_macos.py

필요 패키지:
    pip install coremltools onnx
"""

import coremltools as ct
import numpy as np

# 입력 파일
ONNX_PATH = "model.onnx"  # 같은 폴더에 있어야 함
OUTPUT_PATH = "DriverBehavior.mlpackage"

print("=" * 60)
print("ONNX → CoreML FP16 변환 (macOS)")
print("=" * 60)

# 변환
print("\n[1] CoreML 변환 중...")
mlmodel = ct.convert(
    ONNX_PATH,
    minimum_deployment_target=ct.target.iOS15,
    convert_to="mlprogram",
    compute_precision=ct.precision.FLOAT16,  # FP16 (Metal 최적화)
    compute_units=ct.ComputeUnit.ALL,  # CPU + GPU + Neural Engine
)

# 메타데이터
mlmodel.author = "C-Team"
mlmodel.short_description = "Driver Behavior Detection - Video Swin Transformer"
mlmodel.version = "1.0"

# 저장
mlmodel.save(OUTPUT_PATH)
print(f"✓ 저장 완료: {OUTPUT_PATH}")

# 검증
print("\n[2] 모델 검증 중...")
import onnxruntime as ort

# 더미 입력
dummy_input = np.random.randn(1, 3, 30, 224, 224).astype(np.float32)

# ONNX 출력
sess = ort.InferenceSession(ONNX_PATH)
onnx_out = sess.run(None, {'video_input': dummy_input})[0]

# CoreML 출력
coreml_out = mlmodel.predict({'video_input': dummy_input})['class_logits']

# 비교
max_diff = np.abs(onnx_out - coreml_out).max()
print(f"  ONNX vs CoreML 최대 차이: {max_diff:.6f}")

onnx_pred = onnx_out.argmax()
coreml_pred = coreml_out.argmax()

if onnx_pred == coreml_pred:
    print("  ✓ 예측 클래스 일치!")
else:
    print(f"  ⚠ 예측 다름: ONNX={onnx_pred}, CoreML={coreml_pred}")

print("\n" + "=" * 60)
print("변환 완료!")
print("=" * 60)
print(f"""
iOS 사용법:
1. {OUTPUT_PATH}를 Xcode 프로젝트에 드래그앤드롭
2. 자동으로 Metal FP16 가속 적용됨

입력 형식:
- Shape: [1, 3, 30, 224, 224]
- 정규화: (pixel/255 - [0.485,0.456,0.406]) / [0.229,0.224,0.225]
- 채널: RGB

출력:
- [1, 5] 로짓 → argmax로 클래스 예측
- 0: 정상, 1: 졸음운전, 2: 물건찾기, 3: 휴대폰사용, 4: 운전자폭행
""")