import os
import sys
import onnx

# === 動態加入 optimizer_scripts 模組路徑 ===
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(current_dir, 'tools'))

from optimizer_scripts.pytorch_exported_onnx_preprocess import torch_exported_onnx_flow

def main():
    # === 設定路徑 ===
    onnx_path = r'C:\Users\rd_de\kneronstdc\work_dirs\meconfig8\latest.onnx'
    optimized_path = r'C:\Users\rd_de\kneronstdc\work_dirs\meconfig8\latest_optimized.onnx'

    if not os.path.exists(onnx_path):
        print(f'❌ 找不到 ONNX 檔案: {onnx_path}')
        return

    # === 載入 ONNX 模型 ===
    print(f'🔄 載入 ONNX: {onnx_path}')
    m = onnx.load(onnx_path)

    # === 修正 ir_version（避免 opset11 時報錯）===
    if m.ir_version == 7:
        print(f'⚠️ 調整 ir_version 7 → 6（相容性修正）')
        m.ir_version = 6

    # === 執行 Kneron 優化流程 ===
    print('⚙️ 執行 Kneron 優化 flow...')
    try:
        m = torch_exported_onnx_flow(m, disable_fuse_bn=False)
    except Exception as e:
        print(f'❌ 優化失敗: {type(e).__name__} → {e}')
        return

    # === 儲存結果 ===
    os.makedirs(os.path.dirname(optimized_path), exist_ok=True)
    onnx.save(m, optimized_path)
    print(f'✅ 已儲存最佳化 ONNX: {optimized_path}')

if __name__ == '__main__':
    main()