import ktc
import numpy as np
import os
import onnx
import shutil
from PIL import Image

# === 1. 設定路徑與參數 ===
onnx_dir = 'work_dirs/meconfig8/'
onnx_path = os.path.join(onnx_dir, 'latest.onnx')
optimized_path = os.path.join(onnx_dir, 'latest_optimized.onnx')
data_path = 'data724362'
imgsz_w, imgsz_h = 724, 362  # STDC 預設解析度

# === 2. 建立輸出資料夾 ===
os.makedirs(onnx_dir, exist_ok=True)

# === 3. 優化 ONNX 模型（使用 onnx2onnx_flow）===
print("⚙️ 使用 onnx2onnx_flow 優化 ONNX...")
model = onnx.load(onnx_path)
model = ktc.onnx_optimizer.onnx2onnx_flow(model)
onnx.save(model, optimized_path)

# === 4. 驗證輸入 Shape 是否正確 ===
print("📏 驗證 ONNX Input Shape...")
input_tensor = model.graph.input[0]
input_name = input_tensor.name
input_shape = [d.dim_value for d in input_tensor.type.tensor_type.shape.dim]
expected_shape = [1, 3, imgsz_h, imgsz_w]
print(f"📌 input_name: {input_name}")
print(f"📌 input_shape: {input_shape}")
if input_shape != expected_shape:
    raise ValueError(f"❌ Shape mismatch: {input_shape} ≠ {expected_shape}")

# === 5. 初始化模型編譯器 (for KL630) ===
print("📐 配置模型 for KL630...")
km = ktc.ModelConfig(32769, "0001", "630", onnx_model=model)

# （可選）效能分析
eval_result = km.evaluate()
print("\n📊 NPU 效能分析:\n" + str(eval_result))

# === 6. 圖片預處理 ===
print("🖼️ 處理輸入圖片...")
img_list = []
files_found = [f for _, _, files in os.walk(data_path)
               for f in files if f.lower().endswith((".jpg", ".jpeg", ".png", ".bmp"))]
if not files_found:
    raise FileNotFoundError(f"❌ 找不到圖片於 {data_path}")

for root, _, files in os.walk(data_path):
    for f in files:
        if not f.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
            continue
        fullpath = os.path.join(root, f)
        try:
            img = Image.open(fullpath).convert("RGB")
            img = Image.fromarray(np.array(img)[..., ::-1])  # RGB ➝ BGR
            img_np = np.array(img.resize((imgsz_w, imgsz_h), Image.BILINEAR)).astype(np.float32)
            img_np = img_np / 256.0 - 0.5
            img_np = np.transpose(img_np, (2, 0, 1))  # HWC ➝ CHW
            img_np = np.expand_dims(img_np, axis=0)   # CHW ➝ NCHW
            img_list.append(img_np)
            print(f"✅ 處理成功: {fullpath}")
        except Exception as e:
            print(f"❌ 處理失敗 {fullpath}: {e}")

if not img_list:
    raise RuntimeError("❌ 沒有成功處理任何圖片！")

# === 7. 執行 BIE 量化分析 ===
print("📦 執行固定點分析 (BIE)...")
bie_model_path = km.analysis({input_name: img_list})
bie_save_path = os.path.join(onnx_dir, os.path.basename(bie_model_path))
shutil.copy(bie_model_path, bie_save_path)

if not os.path.exists(bie_save_path):
    raise RuntimeError("❌ 無法產生 BIE 模型")

print("✅ BIE 模型儲存於:", bie_save_path)

# === 8. 編譯 NEF 模型 ===
print("⚙️ 編譯 NEF 模型 for KL630...")
nef_model_path = ktc.compile([km])
nef_save_path = os.path.join(onnx_dir, os.path.basename(nef_model_path))
shutil.copy(nef_model_path, nef_save_path)

if not os.path.exists(nef_save_path):
    raise RuntimeError("❌ 無法產生 NEF 模型")

print("✅ NEF 編譯完成")
print("📁 NEF 檔案儲存於:", nef_save_path)