import os
import numpy as np
import torch
import ktc  # Kneron Toolchain
import onnx
from yolov5_preprocess import Yolov5_preprocess
import kneron_preprocessing

# 設定裝置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 設定圖片大小
imgsz_h, imgsz_w = 640, 640

# 設定正確的數據目錄
data_path = "/workspace/yolov5/data50"

# 確保 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 模型的輸入名稱**
onnx_model_path = "/workspace/yolov5/runs/train/exp24/weights/latest.opt.onnx"
m = onnx.load(onnx_model_path)
input_name = m.graph.input[0].name  # 確保 key 與 ONNX input name 一致

km = ktc.ModelConfig(20008, "0001", "720", onnx_model=onnx_model_path)
# 存儲預處理後的圖片數據
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)