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

# === 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 模型（使用 kneronnxopt API）===
print("⚙️ 使用 kneronnxopt 優化 ONNX...")
try:
    model = onnx.load(onnx_path)
    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]
    print(f"📌 模型實際的 input name 是: {input_name}")

    model = kneronnxopt.optimize(
        model,
        duplicate_shared_weights=1,
        skip_check=False,
        skip_fuse_qkv=True
    )
    onnx.save(model, optimized_path)
except Exception as e:
    print(f"❌ 優化失敗: {e}")
    exit(1)

# === 4. 載入優化後的模型 ===
print("🔄 載入優化後的 ONNX...")
m = onnx.load(optimized_path)

# === 5. 設定 Kneron 模型編譯參數 ===
print("📐 配置模型...")
km = ktc.ModelConfig(20008, "0001", "630", onnx_model=m)

# （可選）模型效能評估
eval_result = km.evaluate()
print("\n📊 NPU 效能評估:\n" + str(eval_result))

# === 6. 處理輸入圖片 ===
print("🖼️ 處理輸入圖片...")
input_name = m.graph.input[0].name
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:
        fullpath = os.path.join(root, f)
        if not f.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
            continue
        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 模型...")
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)