"""Real-geometry image-aware QSA preparation, checked independently with NumPy."""
import hashlib
import json
from pathlib import Path
import subprocess
import sys

checkout, output = map(Path, sys.argv[1:3])
assert subprocess.check_output(["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
sys.path.insert(0, str(checkout))
import mlx.core as mx
import numpy as np
from mtplx.models.qwen4_exp import TextArgs, QSAIndexer, _build_mrope_axes
from mtplx.attention_context import vision_rope

output.mkdir(parents=True, exist_ok=True)
args = TextArgs(hidden_size=640, head_dim=256, indexer_head_dim=128, indexer_n_heads=4,
    indexer_kv_heads=1, indexer_compress_ratio=4, indexer_budget=2048,
    rope_parameters=dict(mrope_interleaved=True, mrope_section=[11,11,10], partial_rotary_factor=.25,
                         rope_theta=10000000, rope_type="default"))
idx = QSAIndexer(args)
axes = _build_mrope_axes([11,11,10],True)
table = np.stack([np.arange(2060, dtype=np.int32)//(axis+1) for axis in range(3)])
def pattern(shape, seed):
    return (((np.arange(np.prod(shape), dtype=np.int64)*17+seed)%97-48)/32).astype(np.float32).reshape(shape)
raw_q = pattern((1,16,4,128),3)
raw_k = pattern((1,2072,128),11)

for dtype_name, dtype in [("f32",mx.float32),("bf16",mx.bfloat16)]:
    idx.q_layernorm.weight = mx.ones((128,),dtype=dtype)
    idx.k_layernorm.weight = mx.ones((128,),dtype=dtype)
    q, k = mx.array(raw_q).astype(dtype), mx.array(raw_k).astype(dtype)
    for mode in ["table", "equal", "delta", "empty"]:
        delta = -7 if mode in ("table","delta","empty") else 0
        table_value = table if mode=="table" else np.broadcast_to(np.arange(2060,dtype=np.int32),(3,2060)) if mode=="equal" else np.zeros((3,0),dtype=np.int32) if mode=="empty" else None
        with vision_rope(None if table_value is None else mx.array(table_value),delta):
            query = idx._prepare_queries(q,2056)
            pooled = idx._pool_keys_eager(k,0,518)
        for name,value in [("query",query),("pooled",pooled)]:
            data=np.asarray(value.astype(mx.float32))
            (output/f"{dtype_name}-{mode}-{name}.f32").write_bytes(data.tobytes())
            print(json.dumps(dict(case="indexer",dtype=dtype_name,mode=mode,part=name,shape=list(data.shape),sha256=hashlib.sha256(data.tobytes()).hexdigest())),flush=True)
        if dtype_name=="f32":
            def independent(values, positions):
                values = values / np.sqrt(np.mean(values*values,axis=-1,keepdims=True)+args.rms_norm_eps)
                pos3=np.stack([table_value[:,p] if table_value is not None and p<table_value.shape[1] else np.full(3,p+delta) for p in positions],axis=1)
                angles=pos3[axes].T.astype(np.float32)*np.asarray(idx._inv_freq)[None,:]
                angles=np.concatenate([angles,angles],axis=-1)[None,:,None,:]
                first=values[...,:64];rotated=np.concatenate([-first[...,32:],first[...,:32]],axis=-1)
                return np.concatenate([first*np.cos(angles)+rotated*np.sin(angles),values[...,64:]],axis=-1)
            np.testing.assert_allclose(np.asarray(query),independent(raw_q,np.arange(2056,2072)),rtol=3e-6,atol=3e-6)
            np.testing.assert_allclose(np.asarray(pooled),independent(raw_k.reshape(1,518,4,128).mean(axis=2)[:,:,None],np.arange(0,2072,4))[:,:,0],rtol=3e-6,atol=3e-6)
print(json.dumps(dict(case="independent_numpy",passed=True)),flush=True)