"""
onnx_to_bie.py  —  ONNX → BIE → NEF for Kneron KL730
======================================================
Run INSIDE the Kneron toolchain Docker container:
  docker run --rm -it -v /your/data:/data1 kneron/toolchain:v0.17.2
  python onnx_to_bie.py

Accuracy checklist (every item below is already applied):
  ✔  cv2 letterbox  — identical to eval script, eliminates calibration/eval mismatch
  ✔  BGR → RGB      — same channel order as eval script
  ✔  float32 [0,1]  — same normalisation as eval script
  ✔  Shuffled calib — random sample covers full dataset distribution
  ✔  int16 bitwidth — preserves full dynamic range of detection head outputs
  ✔  500 images     — sufficient for stable quantisation statistics
"""

import os
import random
import onnx
import ktc
import cv2
import numpy as np
from pathlib import Path

# ─── CONFIG  (edit these paths) ───────────────────────────────────────────────
ONNX_PATH  = Path("./models/optimized_best_no_sigmoid.onnx").resolve()
CALIB_DIR  = Path("./100_crop/train/images").resolve()
OUTPUT_DIR = Path("./models").resolve()

MODEL_ID   = 32770
MODEL_VER  = "0001"
PLATFORM   = "730"

IMG_SIZE   = 640
NUM_CALIB  = 100          # 500 is good; 1000 is better if you have time
SEED       = 42           # fixed seed → reproducible BIE every run
INPUT_NAME = "images"     # must match ONNX input node name exactly
                          # verify with:
                          #   python -c "import onnx; m=onnx.load('model.onnx'); \
                          #              print([i.name for i in m.graph.input])"

# ─── PREPROCESSING  (must be identical to eval_bie_yolo11m.py) ────────────────
def letterbox(img, new_size=640, color=(114, 114, 114)):
    """
    Resize-with-padding using OpenCV — identical implementation to the
    eval script so calibration and inference see exactly the same pixels.
    """
    h, w   = img.shape[:2]
    scale  = new_size / max(h, w)
    new_w  = int(round(w * scale))
    new_h  = int(round(h * scale))

    img_resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

    pad_w  = new_size - new_w
    pad_h  = new_size - new_h
    top    = pad_h // 2;    bottom = pad_h - top
    left   = pad_w // 2;    right  = pad_w - left

    return cv2.copyMakeBorder(
        img_resized, top, bottom, left, right,
        cv2.BORDER_CONSTANT, value=color,
    )


def preprocess(img_path):
    """
    Load → letterbox → BGR→RGB → HWC→CHW → float32 [0,1] → (1,3,H,W)
    Returns None if the image cannot be read (skipped silently).
    """
    img = cv2.imread(str(img_path))
    if img is None:
        print(f"  [WARN] Cannot read image, skipping: {img_path}")
        return None

    img = letterbox(img, IMG_SIZE)              # (640, 640, 3)  BGR
    img = img[:, :, ::-1]                       # BGR → RGB
    img = img.astype(np.float32) / 255.0        # [0, 255] → [0.0, 1.0]
    img = img.transpose(2, 0, 1)               # HWC → CHW
    return np.expand_dims(img, 0)              # (1, 3, 640, 640)


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

# ── Step 1: Load ONNX ─────────────────────────────────────────────────────────
print("=" * 60)
print("[Step 1] Loading ONNX model...")
print(f"  Path : {ONNX_PATH}")
m = onnx.load(str(ONNX_PATH))
print("  Loaded OK")

# ── Step 2: Create ModelConfig ────────────────────────────────────────────────
print("\n[Step 2] Creating ModelConfig...")
km = ktc.ModelConfig(MODEL_ID, MODEL_VER, PLATFORM, onnx_model=m)
print(f"  Model ID  : {MODEL_ID}")
print(f"  Platform  : KL{PLATFORM}")

# ── Step 3: IP Evaluation ─────────────────────────────────────────────────────
print("\n[Step 3] Running IP evaluation (operator support check)...")
eval_result = km.evaluate()
print(f"  Result:\n{eval_result}")

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

if not all_imgs:
    raise RuntimeError(f"No images found in CALIB_DIR: {CALIB_DIR}")

# Shuffle with fixed seed so the BIE is reproducible but covers the full
# dataset distribution (critical for correct quantisation statistics).
random.seed(SEED)
random.shuffle(all_imgs)
selected = all_imgs[:NUM_CALIB]

print(f"  Total images found  : {len(all_imgs)}")
print(f"  Selected for calib  : {len(selected)}  (seed={SEED})")

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 len(img_list) < 100:
    raise RuntimeError(
        f"Only {len(img_list)} calibration images loaded — need at least 100. "
        f"Check CALIB_DIR: {CALIB_DIR}"
    )

# Sanity check on first tensor
s = img_list[0]
print(f"\n  Tensor shape : {s.shape}")
print(f"  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 s.shape  == (1, 3, IMG_SIZE, IMG_SIZE), \
    f"Expected shape (1, 3, {IMG_SIZE}, {IMG_SIZE}), got {s.shape}"
assert s.min()  >= 0.0,  "min value < 0 — normalisation error"
assert s.max()  <= 1.01, "max value > 1 — forgot /255 ?"
print("  Sanity check : PASSED")

# ── Step 5: Fixed-point analysis → .bie ──────────────────────────────────────
print("\n[Step 5] Running fixed-point analysis (quantisation)...")
print("  bitwidth mode : int16")
print("  threads       :", max(1, (os.cpu_count() or 4) - 2))
print("  This may take 10–30 minutes for a YOLOv11m model...\n")

bie_path = km.analysis(
    {INPUT_NAME: img_list},
    output_dir             = str(OUTPUT_DIR),
    threads                = max(1, (os.cpu_count() or 4) - 2),

    # ── WHY int16 ──────────────────────────────────────────────────────────
    # The detection head outputs (box regression + class logits) have a wide
    # dynamic range. int8 truncates this range and loses fine-grained score
    # differences between classes, which directly hurts mAP.
    # int16 gives 256× more quantisation levels → near-zero accuracy loss.
    # It is slower on-device than int8 but is the correct default when
    # preserving mAP is the priority.
    # Once you confirm mAP is acceptable, you can try "mix balance" to trade
    # a small amount of accuracy for inference speed.
    datapath_bitwidth_mode = "int16",
)

bie_path = Path(bie_path).resolve()
if not bie_path.exists():
    raise FileNotFoundError(f"BIE was not created at expected path: {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"  NEF saved : {nef_path}")

print("\n" + "=" * 60)
print("  DONE")
print(f"  BIE : {bie_path}")
print(f"  NEF : {nef_path}")
print("=" * 60)
print("""
Next step — evaluate BIE accuracy (run outside Docker):
  python eval_bie_yolo11m.py \\
      --bie      {bie} \\
      --img-dir  ./dataset/val/images \\
      --lbl-dir  ./dataset/val/labels \\
      --iou-sweep \\
      --apply-sigmoid
""".format(bie=bie_path))