#!/usr/bin/env python
"""
Step 3: build the calibration NPZ (ModelOpt layout) and the evaluation NPZ for the vision encoder at a fixed
H x W input (default 392 x 700 -> grid_thw [1, 28, 50] -> 1400 patch tokens -> 350 output tokens).

Images (all already in the local HF cache, no download):
  calibration : lmms-lab/MMMU 'dev' split (the set TensorRT-Edge-LLM itself calibrates vision encoders on)
                interleaved with COCO-2017 val natural photos (detection-datasets/coco, file 0)
  evaluation  : the 4 sample pictures shipped with TensorRT-Edge-LLM + MMMU 'validation' split
                + COCO-2017 val photos from file 1  (all disjoint from the calibration images)

Pre-processing = the official Qwen2VLImageProcessor (bicubic resize, rescale, CLIP mean/std normalise,
temporal duplication, 14x14 patchify in 2x2-merge order) after resizing each image to exactly W x H.
The remaining ONNX inputs (rotary_pos_emb, window_index, cu_window_seqlens, ...) are produced by the
Hugging Face Qwen2.5-VL vision module's own rot_pos_emb() / get_window_index() so they match the exporter.

ModelOpt NPZ layout : one array per graph input, samples concatenated along axis 0; ModelOpt splits every
                      input into N chunks using dim0 of --calibration_shapes (fixed-size samples required).
Eval NPZ layout     : one array per input with a leading sample dim, plus 'ids'.
"""
import argparse, glob, io, json, os
import numpy as np
import pyarrow.parquet as pq
import torch
from PIL import Image
from transformers import Qwen2VLImageProcessor
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VisionTransformerPretrainedModel

ROOT = os.environ.get("ROOT", "/data/users/logesh/Infernece_vision_Manual")
HF = os.path.expanduser("~/.cache/huggingface/hub")
ap = argparse.ArgumentParser()
ap.add_argument("--height", type=int, default=int(os.environ.get("IMG_H", 392)))
ap.add_argument("--width", type=int, default=int(os.environ.get("IMG_W", 700)))
ap.add_argument("--n_calib", type=int, default=128)
ap.add_argument("--n_eval", type=int, default=100)
ap.add_argument("--out_dir", default=f"{ROOT}/data")
ap.add_argument("--proc_dir", default=os.environ.get("SRC_VISUAL_DIR", "/data/users/logesh/TensorRT-Edge-LLM/Qwen/Qwen3-VL-2B-Instruct/onnx/visual"))
ap.add_argument("--min_side", type=int, default=64, help="skip tiny source images (icons)")
args = ap.parse_args()

cfg = json.load(open(os.path.join(args.proc_dir, "config.json")))["vision_config"]
PATCH, MERGE = cfg["patch_size"], cfg["spatial_merge_size"]
assert args.height % (PATCH * MERGE) == 0 and args.width % (PATCH * MERGE) == 0, "H and W must be multiples of 28"
GRID = (1, args.height // PATCH, args.width // PATCH)

# tiny (depth=1, random weights) HF vision module: only used for rot_pos_emb / get_window_index (weight-free)
vcfg = Qwen2_5_VLVisionConfig(**{k: v for k, v in cfg.items() if k in Qwen2_5_VLVisionConfig().to_dict()})
vcfg.depth = 1
vcfg._attn_implementation = "eager"
helper = Qwen2_5_VisionTransformerPretrainedModel(vcfg).float().eval()
proc = Qwen2VLImageProcessor.from_pretrained(args.proc_dir)

def make_sample(img: Image.Image) -> dict:
    img = img.convert("RGB").resize((args.width, args.height), Image.BICUBIC)
    out = proc(images=[img], return_tensors="np")
    pv, thw = out["pixel_values"], out["image_grid_thw"]
    assert tuple(int(x) for x in thw[0]) == GRID, (thw, GRID)
    hw = pv.shape[0]
    thw_t = torch.as_tensor(thw, dtype=torch.int64)
    with torch.no_grad():
        rpe = helper.rot_pos_emb(thw_t).float().numpy()                     # [hw, 40] fp32
        wi, cuw = helper.get_window_index(thw_t)
    cuw = torch.unique_consecutive(torch.tensor(cuw, dtype=torch.int32)).numpy()
    wi = wi.to(torch.int64).numpy()
    return {
        "input": pv.astype(np.float16),                        # [hw, 1176]
        "rotary_pos_emb": rpe.astype(np.float32),              # [hw, 40]
        "cu_seqlens": np.array([0, hw], dtype=np.int32),       # [2]
        "max_seqlen_carrier": np.zeros(hw, dtype=np.int32),    # [hw]  shape-only carrier
        "cu_window_seqlens": cuw.astype(np.int32),             # [n_windows + 1]
        "window_index": wi,                                    # [hw/4] int64
        "reverse_window_index": np.argsort(wi).astype(np.int64),
    }

def ok(im):  # skip icons / degenerate images
    return min(im.size) >= args.min_side

def mmmu_images(split, limit, skip=0):
    f = glob.glob(f"{HF}/datasets--lmms-lab--MMMU/snapshots/*/data/{split}-00000-of-00001.parquet")[0]
    rows = pq.read_table(f, columns=["id"] + [f"image_{i}" for i in range(1, 8)]).to_pylist()
    imgs = []
    for row in rows:
        for i in range(1, 8):
            v = row[f"image_{i}"]
            if v and v.get("bytes"):
                im = Image.open(io.BytesIO(v["bytes"]))
                if ok(im):
                    imgs.append((f"mmmu-{split}:{row['id']}#img{i}", im))
                if len(imgs) >= limit + skip:
                    return imgs[skip:]
    return imgs[skip:]

def coco_images(file_idx, limit):
    files = sorted(glob.glob(f"{HF}/datasets--detection-datasets--coco/snapshots/*/data/val-*.parquet"))
    pf = pq.ParquetFile(files[file_idx])
    imgs = []
    for batch in pf.iter_batches(batch_size=64, columns=["image_id", "image"]):
        for row in batch.to_pylist():
            im = Image.open(io.BytesIO(row["image"]["bytes"]))
            if ok(im):
                imgs.append((f"coco-val:{row['image_id']}", im))
            if len(imgs) >= limit:
                return imgs
    return imgs

def interleave(a, b):
    out = []
    for i in range(max(len(a), len(b))):
        if i < len(a): out.append(a[i])
        if i < len(b): out.append(b[i])
    return out

os.makedirs(args.out_dir, exist_ok=True)
tag = f"{args.height}x{args.width}"

# ------------------------------------------------------------------ calibration (ModelOpt layout)
half = (args.n_calib + 1) // 2
calib = interleave(mmmu_images("dev", half), coco_images(0, args.n_calib - half))[:args.n_calib]
samples = [make_sample(im) for _, im in calib]
keys = list(samples[0].keys())
calib_npz = {k: np.concatenate([s[k] for s in samples], 0) for k in keys}   # exactly the 7 graph inputs
calib_path = os.path.join(args.out_dir, f"calib_{tag}.npz")
np.savez(calib_path, **calib_npz)
shapes = ",".join(f"{k}:{'x'.join(str(d) for d in samples[0][k].shape)}" for k in keys)
open(os.path.join(args.out_dir, f"calib_{tag}.shapes"), "w").write(shapes + "\n")
open(os.path.join(args.out_dir, f"calib_{tag}.ids.txt"), "w").write("\n".join(n for n, _ in calib) + "\n")
print(f"calibration: {len(samples)} images -> {calib_path}")
print("  --calibration_shapes", shapes)
print("  ", {k: (v.shape, str(v.dtype)) for k, v in calib_npz.items()})
pv = calib_npz["input"].astype(np.float32)
print(f"  pixel_values stats: min {pv.min():.3f} max {pv.max():.3f} mean {pv.mean():.4f} std {pv.std():.4f} nan {np.isnan(pv).any()}")

# ------------------------------------------------------------------ evaluation (disjoint images)
ev = [(f"edgellm:{os.path.basename(p)}", Image.open(p))
      for p in sorted(glob.glob("/data/users/logesh/TensorRT-Edge-LLM/examples/multimodal/pics/*.jpeg"))]
rest = args.n_eval - len(ev)
ev += interleave(mmmu_images("validation", (rest + 1) // 2), coco_images(1, rest // 2))[:rest]
samples = [make_sample(im) for _, im in ev]
eval_npz = {k: np.stack([s[k] for s in samples], 0) for k in keys}
eval_npz["ids"] = np.array([n for n, _ in ev])
eval_path = os.path.join(args.out_dir, f"eval_{tag}.npz")
np.savez(eval_path, **eval_npz)
print(f"evaluation: {len(samples)} images -> {eval_path}")
print("  ", {k: (v.shape, str(v.dtype)) for k, v in eval_npz.items()})
assert not (set(n for n, _ in calib) & set(n for n, _ in ev)), "calibration / evaluation overlap"

# ------------------------------------------------------------------ raw inputs for trtexec --loadInputs (sample 0)
raw_dir = os.path.join(args.out_dir, f"trtexec_inputs_{tag}")
os.makedirs(raw_dir, exist_ok=True)
spec = []
for k in keys:
    p = os.path.join(raw_dir, f"{k}.bin")
    np.ascontiguousarray(samples[0][k]).tofile(p)
    spec.append(f"{k}:{p}")
open(os.path.join(raw_dir, "loadInputs.txt"), "w").write(",".join(spec) + "\n")
print("trtexec --loadInputs spec written to", os.path.join(raw_dir, "loadInputs.txt"))