"""Qwen3.5/3.6 vision: numpy patchify vs the HF processor, and the static OV tower vs the HF tower (CPU).
usage: python check_vision.py HF_DIR IMAGE [DEVICE=CPU]"""
import json, os, sys
import numpy as np
import openvino as ov
import torch
from PIL import Image
from transformers import AutoProcessor

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from onw.models import qwen_vision as QV
from onw.models.lfm2_moe import Checkpoint


def main():
    src, img_path = sys.argv[1], sys.argv[2]
    dev = sys.argv[3] if len(sys.argv) > 3 else "CPU"
    grid = (32, 32)
    img = Image.open(img_path).convert("RGB").resize((512, 512), Image.BICUBIC)
    vcfg = {k: v for k, v in json.load(open(os.path.join(src, "config.json")))["vision_config"].items() if k != "model_type"}
    proc = AutoProcessor.from_pretrained(src)
    ref = proc.image_processor(images=[img], return_tensors="pt")
    ref_px, thw = ref["pixel_values"].numpy(), ref["image_grid_thw"]
    pp = json.load(open(os.path.join(src, "preprocessor_config.json")))
    mine = QV.patchify(img, grid, mean=pp["image_mean"], std=pp["image_std"])
    print("grid", thw.tolist(), "patchify rel err", float(np.linalg.norm(mine - ref_px) / np.linalg.norm(ref_px)))
    from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel
    from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5VisionConfig
    hf = Qwen3_5VisionModel._from_config(Qwen3_5VisionConfig(**vcfg))
    W = Checkpoint(src)
    sd = {k[len(QV.VP):]: torch.tensor(W[k]) for k in W.where if k.startswith(QV.VP)}
    print("load:", hf.load_state_dict(sd, strict=False))
    with torch.no_grad():
        want = hf(torch.tensor(ref_px), thw).pooler_output.float().numpy()
    m = QV.build(src, vcfg, W, grid)
    got = ov.Core().compile_model(m, dev, {"INFERENCE_PRECISION_HINT": "f32"} if dev == "CPU" else {}).create_infer_request() \
        .infer([mine.astype(np.float16)])[0].astype(np.float32)
    print(f"{dev} tower vs HF: shape {got.shape} vs {want.shape}, rel err {np.linalg.norm(got - want) / np.linalg.norm(want):.4f}")


if __name__ == "__main__":
    main()