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

# 設定 ONNX 模型儲存路徑
onnx_dir = 'runs/train/exp73/weights/'
onnx_path = os.path.join(onnx_dir, 'best_simplified.onnx')

# 確保目標資料夾存在
os.makedirs(onnx_dir, exist_ok=True)

# 加載並優化 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)

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

print(f"\u2705 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"\u26a0\ufe0f 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"\u2705 Processed: {fullpath}")
            img_list.append(img_data)
        except Exception as e:
            print(f"\u274c Failed to process {fullpath}: {e}")

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

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

# 確保 BIE 檔案儲存到指定目錄
bie_save_path = os.path.join(onnx_dir, os.path.basename(bie_model_path))
shutil.copy(bie_model_path, bie_save_path)  # 使用 shutil.move 來處理跨磁碟移動

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

print("\n\u2705 Fixed-point analysis done! BIE model saved to:", bie_save_path)

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

# 確保 nef_model_path 不是 None 或空值
if not nef_model_path:
    raise RuntimeError("❌ Error: ktc.compile() did not return a valid .nef file path!")

# 確保 NEF 目標資料夾存在
os.makedirs(onnx_dir, exist_ok=True)

# 確保 nef_model_path 不是 None 或空值
if not nef_model_path:
    raise RuntimeError("❌ Error: ktc.compile() did not return a valid .nef file path!")

# 確保 .nef 檔案存在
if not os.path.exists(nef_model_path):
    raise RuntimeError(f"❌ Error: NEF model was not generated at {nef_model_path}! Please check your compilation process.")

# 確保 NEF 檔案儲存到指定目錄
nef_save_path = os.path.join(onnx_dir, os.path.basename(nef_model_path))
if os.path.exists(nef_model_path):
    shutil.copy(nef_model_path, nef_save_path)
else:
    raise RuntimeError(f"❌ Error: NEF model was expected at {nef_model_path}, but it does not exist!")  # 同樣使用 shutil.move

if not os.path.exists(nef_save_path):
    raise RuntimeError(f"\u274c Error: NEF model was not generated! Please check your compilation process.")

print("\n\u2705 Compile done! NEF file saved to:", nef_save_path)