import os
import numpy as np
import onnx
import shutil
import cv2
import ktc

onnx_dir = 'work_dirs/kn_stdc1_in1k-pre_512x1024_80k_cityscapes/'
onnx_path = os.path.join(onnx_dir, 'latest.onnx')
data_path = "data512"
imgsz = (512, 512)

os.makedirs(onnx_dir, exist_ok=True)

print("🔄 Loading and optimizing ONNX...")
model = onnx.load(onnx_path)
model = ktc.onnx_optimizer.onnx2onnx_flow(model)
opt_onnx_path = os.path.join(onnx_dir, 'latest.opt.onnx')
onnx.save(model, opt_onnx_path)

print("📐 Configuring model...")
km = ktc.ModelConfig(20008, "0001", "630", onnx_model=model)

# Optional: performance check
print("\n📊 Evaluating model...")
print(km.evaluate())

input_name = model.graph.input[0].name
print("📥 ONNX input name:", input_name)

img_list = []
print("🖼️ Preprocessing images...")
for root, _, files in os.walk(data_path):
    for fname in files:
        if fname.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            path = os.path.join(root, fname)
            img = cv2.imread(path)
            img = cv2.resize(img, imgsz)
            img = img.astype(np.float32) / 256.0 - 0.5
            img = np.transpose(img, (2, 0, 1))  # HWC ➝ CHW
            img = np.expand_dims(img, axis=0)   # Add batch dim
            img_list.append(img)
            print("✅", path)

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

print("📦 Quantizing (BIE)...")
bie_path = km.analysis({input_name: img_list})
bie_save = os.path.join(onnx_dir, os.path.basename(bie_path))
shutil.copy(bie_path, bie_save)

if not os.path.exists(bie_save):
    raise RuntimeError("❌ BIE model not saved!")

print("⚙️ Compiling NEF...")
nef_path = ktc.compile([km])
nef_save = os.path.join(onnx_dir, os.path.basename(nef_path))
shutil.copy(nef_path, nef_save)

if not os.path.exists(nef_save):
    raise RuntimeError("❌ NEF model not saved!")

print("✅ Compile finished. NEF at:", nef_save)