"""
onnx_to_bie.py  —  Simple ONNX → BIE conversion for Kneron KL730
Run inside the Kneron toolchain docker container:
  docker run --rm -it -v /your/data:/data1 kneron/toolchain:v0.17.2
"""

import os
import onnx
import ktc
import numpy as np
from pathlib import Path
from PIL import Image
# ─── CONFIG ───────────────────────────────────────────────────────────────────
ONNX_PATH    = Path("./models/optimized_best_no_sigmoid.onnx").resolve()
CALIB_DIR    = Path("./500_crop/train/images").resolve()
OUTPUT_DIR   = Path("./models").resolve()

MODEL_ID     = 32770
MODEL_VER    = "0001"      # version string, keep as "0001" if unsure
PLATFORM     = "730"

IMG_SIZE     = 640
NUM_CALIB    = 500         # number of calibration images

# ─── PREPROCESS ───────────────────────────────────────────────────────────────
# def letterbox(img, size=640):
#     h, w   = img.shape[:2]
#     scale  = size / max(h, w)
#     new_w  = int(round(w * scale))
#     new_h  = int(round(h * scale))
#     img    = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
#     pad_w  = size - new_w
#     pad_h  = size - new_h
#     top    = pad_h // 2;  bottom = pad_h - top
#     left   = pad_w // 2;  right  = pad_w - left
#     return cv2.copyMakeBorder(img, top, bottom, left, right,
#                                cv2.BORDER_CONSTANT, value=(114, 114, 114))

# def preprocess(img_path):
#     img = cv2.imread(str(img_path))
#     if img is None:
#         return None
#     img = letterbox(img, IMG_SIZE)                   # resize + pad
#     img = img[:, :, ::-1]                            # BGR → RGB
#     img = img.astype(np.float32) / 255.0             # normalize [0, 1]
#     img = img.transpose(2, 0, 1)                     # HWC → CHW
#     return np.expand_dims(img, axis=0)               # (1, 3, 640, 640)

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114)):
    """Letterbox resize maintaining aspect ratio"""
    shape = im.size[::-1]  # PIL gives (W, H), we need (H, W)
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
    
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]
    dw /= 2
    dh /= 2
    
    if shape[::-1] != new_unpad:
        im = im.resize(new_unpad, Image.BILINEAR)
    
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    
    new_im = Image.new('RGB', (new_shape[1], new_shape[0]), color)
    new_im.paste(im, (left, top))
    return new_im

def preprocess(img_path):
    """Preprocess with letterbox"""
    im = Image.open(img_path).convert("RGB")
    im = letterbox(im, (640, 640))
    im = np.array(im).astype(np.float32) / 255.0
    im = np.transpose(im, (2, 0, 1))
    return np.expand_dims(im, 0)


# ─── MAIN ─────────────────────────────────────────────────────────────────────
os.makedirs(OUTPUT_DIR, exist_ok=True)

# # ── Step 1: Load and optimize ONNX ───────────────────────────────────────────
# print("=" * 60)
# print("[Step 1] Loading and optimizing ONNX...")
m = onnx.load(str(ONNX_PATH))
# m = ktc.onnx_optimizer.onnx2onnx_flow(m)
# opt_path = ONNX_PATH.replace(".onnx", ".opt.onnx")
# onnx.save(m, opt_path)
# print(f"  Optimized ONNX saved: {opt_path}")

# ── Step 2: Create ModelConfig ────────────────────────────────────────────────
print("\n[Step 2] Creating ModelConfig...")
km = ktc.ModelConfig(MODEL_ID, MODEL_VER, PLATFORM, onnx_model=m)

# ── Step 3: IP Evaluation (optional but useful) ───────────────────────────────
print("\n[Step 3] Running IP evaluation...")
eval_result = km.evaluate()
print(f"  NPU eval result:\n{eval_result}")

# ── Step 4: Load calibration images ──────────────────────────────────────────
print(f"\n[Step 4] Loading calibration images from {CALIB_DIR}...")
exts = (".jpg", ".jpeg", ".png", ".bmp")
all_imgs = sorted([
    f for f in Path(CALIB_DIR).iterdir()
    if f.suffix.lower() in exts
])

# evenly sample NUM_CALIB images
step = max(1, len(all_imgs) // NUM_CALIB)
selected = all_imgs[::step][:NUM_CALIB]
print(f"  Total images found : {len(all_imgs)}")
print(f"  Selected for calib : {len(selected)}")

img_list = []
for p in selected:
    t = preprocess(p)
    if t is not None:
        img_list.append(t)

print(f"  Successfully loaded: {len(img_list)}")
if not img_list:
    raise RuntimeError("No calibration images loaded — check CALIB_DIR")

# Quick sanity check
s = img_list[0]
print(f"  Tensor shape : {s.shape}  dtype : {s.dtype}")
print(f"  Value range  : min={s.min():.4f}  max={s.max():.4f}")
assert s.dtype == np.float32,      "dtype must be float32"
assert 0.0 <= s.min(),             "min value < 0"
assert s.max() <= 1.01,            "max value > 1 — forgot /255?"

# ── Step 5: Fixed-point analysis → generate .bie ─────────────────────────────
print("\n[Step 5] Running fixed-point analysis (quantization)...")
print("  This may take several minutes...")

# input_mapping key must match the ONNX input node name exactly.
# Run: python -c "import onnx; m=onnx.load('./models/optimized_best_no_sigmoid.onnx'); print([i.name for i in m.graph.input])"
# to verify. Common names: "images", "input", "input.1"
INPUT_NAME = "images"

bie_path = km.analysis(
    {INPUT_NAME: img_list},
    output_dir             = str(OUTPUT_DIR),
    threads                = max(1, (os.cpu_count() or 4) - 2),
    datapath_bitwidth_mode = "mix balance",     # "int8" / "int16" / "mix balance" / "mix light"
    # quantize_mode          = "default",  # "default" — sigmoid already removed from ONNX
)

bie_path = Path(bie_path).resolve()
if not bie_path.exists():
    raise FileNotFoundError(f"BIE file was not created: {bie_path}")

km.bie_path = str(bie_path)

# print(f"\n  BIE saved: {bie_path}")

# ── Step 6: Compile BIE → NEF ────────────────────────────────────────────────
print("\n[Step 6] Compiling BIE → NEF...")
nef_path = ktc.compile([km], output_dir=str(OUTPUT_DIR))
print(f"\n  NEF saved: {nef_path}")

print("\n" + "=" * 60)
print("  Done!")
# print(f"  BIE : {bie_path}")
print(f"  NEF : {nef_path}")
print("=" * 60)


# # First fix the \n bug in the version file
# import ktc

# # Use the BIE that was ALREADY generated successfully
# bie = './models/input.kdp730.scaled.bie'

# # Create a NEW ModelConfig from the BIE — not from onnx
# km_bie = ktc.ModelConfig(32770, '0001', '730', bie_path=bie)

# # Compile
# nef = ktc.compile([km_bie])
# print('NEF:', nef)
