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')
data_path = "data724362"
imgsz_w, imgsz_h = 724, 362  # STDC 預設解析度

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

# === 3. 載入並優化 ONNX 模型 ===
print("🔄 Loading and optimizing ONNX...")
m = onnx.load(onnx_path)
m = ktc.onnx_optimizer.onnx2onnx_flow(m)
opt_onnx_path = os.path.join(onnx_dir, 'latest.opt.onnx')
onnx.save(m, opt_onnx_path)

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

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

# === 5. 準備圖片資料 ===
print("🖼️ Preparing image data...")
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"❌ No images found in {data_path}!")

print(f"✅ Found {len(files_found)} images in {data_path}")

input_name = m.graph.input[0].name
img_list = []

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 (加上 batch 維度)
            img_list.append(img_np)
            print(f"✅ Processed: {fullpath}")
        except Exception as e:
            print(f"❌ Failed to process {fullpath}: {e}")

if not img_list:
    raise RuntimeError("❌ Error: No valid images were processed!")

# === 6. BIE 量化分析 ===
print("📦 Running fixed-point analysis...")
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("❌ Error: BIE model was not generated!")

print("✅ BIE model saved to:", bie_save_path)

# === 7. 編譯 NEF 模型 ===
print("⚙️ Compiling NEF model...")
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("❌ Error: NEF model was not generated!")

print("✅ NEF compile done!")
print("📁 NEF file saved to:", nef_save_path)