import ktc
import numpy as np
import os
import onnx
from PIL import Image
import torch
from yolov5_preprocess import Yolov5_preprocess
import kneron_preprocessing

onnx_path = 'runs/train/exp24/weights/best_simplified.onnx'
m = onnx.load(onnx_path)
m = ktc.onnx_optimizer.onnx2onnx_flow(m)
onnx.save(m,'latest.opt.onnx')
km = ktc.ModelConfig(20008, "0001", "720", onnx_model=m)
eval_result = km.evaluate()
print("\nNpu performance evaluation result:\n" + str(eval_result))

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

imgsz_h, imgsz_w = 640, 640

data_path = "data50"
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"❌ Error: No images found in {data_path}! Please check your dataset.")

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

# **獲取 ONNX 模型的輸入名稱**
input_name = m.graph.input[0].name  # 確保 key 與 ONNX input name 一致
# 存儲預處理後的圖片數據
img_list = []

# 遍歷 data50 並進行預處理
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")):
            print(f"⚠️ Skipping non-image file: {fullpath}")
            continue

        # **嘗試處理圖片**
        try:
            img_data, _ = Yolov5_preprocess(fullpath, device, imgsz_h, imgsz_w)
            img_data = img_data.cpu().numpy()
            print(f"✅ Processed: {fullpath}")
            img_list.append(img_data)
        except Exception as e:
            print(f"❌ Failed to process {fullpath}: {e}")

# **確保 img_list 不是空的**
if not img_list:
    raise ValueError("❌ Error: No valid images were processed! Please check the image paths and formats.")

# **執行 BIE 量化**
bie_model_path = km.analysis({input_name: img_list})

# **確認 BIE 模型是否生成**
if not os.path.exists(bie_model_path):
    raise RuntimeError(f"❌ Error: BIE model was not generated! Please check your quantization process.")

# 顯示成功訊息
print("\n✅ Fixed-point analysis done! BIE model saved to:", bie_model_path)

# 確保 `km` 已經初始化，並且 `.bie` 模型已生成
nef_model_path = ktc.compile([km])

# 顯示成功訊息
print("\n✅ Compile done! NEF file saved to:", nef_model_path)