"""Pinned attention chunk RoPE oracle, with independent per-row controls."""
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 _vision_chunk_cos_sin, _apply_partial_rope, _build_mrope_axes, _rope_cos_sin, _rope_inv_freq_and_scaling, TextArgs

output.mkdir(parents=True, exist_ok=True)
manifest = []
for dim, rotary, sections, heads in [(8, 8, [2,1,1], 2), (256,64,[11,11,10],24)]:
    axes = mx.array(_build_mrope_axes(sections, True), dtype=mx.int32)
    freq, _ = _rope_inv_freq_and_scaling(TextArgs(head_dim=dim, rope_parameters=dict(rope_type="default", rope_theta=10000000, partial_rotary_factor=rotary/dim)))
    short = [[0,1,2,2],[0,1,3,3],[0,1,4,4]]
    full = [r + [3,4,5,6] for r in short]
    equal = [list(range(4))] * 3
    for mode, table, start, rows, delta in [
        ('inside',short,0,4,-1), ('outside',short,4,2,-1),
        ('straddle',short,2,4,-1), ('full_text_tail',full,6,4,-1),
        ('equal',equal,2,4,0), ('none',None,4,2,-1), ('empty',[[],[],[]],4,2,-1),
    ]:
        for scaling in [1.,1.25]:
            t = None if table is None else mx.array(table,dtype=mx.int32)
            cos, sin = _vision_chunk_cos_sin(t,delta,start,rows,freq,axes,scaling)
            # Independent row-position construction; no chunk helper is used.
            pos3 = np.stack([np.array(table)[:,p] if table is not None and p<len(table[0]) else np.full(3,p+delta) for p in range(start,start+rows)],axis=1)
            angles = pos3[np.asarray(axes)].T.astype(np.float32) * np.asarray(freq)[None,:]
            angles = np.concatenate([angles,angles],axis=1)
            np.testing.assert_allclose(np.asarray(cos),np.cos(angles)*scaling,rtol=3e-6,atol=3e-6)
            np.testing.assert_allclose(np.asarray(sin),np.sin(angles)*scaling,rtol=3e-6,atol=3e-6)
            if mode in ('full_text_tail','equal'):
                plain = _rope_cos_sin(mx.arange(start+delta,start+rows+delta,dtype=mx.int32),freq,scaling)
                assert all(mx.array_equal(a,b).item() for a,b in zip((cos,sin),plain))
            for dtype_name,dtype in [('f32',mx.float32),('bf16',mx.bfloat16)]:
                key = f'{dim}-{mode}-{scaling}-{dtype_name}'
                record = dict(key=key,dim=dim,rotary=rotary,sections=sections,heads=heads,mode=mode,table=table,start=start,rows=rows,delta=delta,scaling=scaling,dtype=dtype_name)
                arrays = [('cos',cos),('sin',sin)]
                for part,h,seed in [('q',heads,3),('k',1,11)]:
                    shape = (1,rows,h,dim)
                    raw = (((np.arange(np.prod(shape),dtype=np.int64)*17+seed)%97-48)/32).astype(np.float32).reshape(shape)
                    value = _apply_partial_rope(mx.array(raw).astype(dtype),cos,sin)
                    arrays.append((part,value))
                for part,value in arrays:
                    data=np.asarray(value.astype(mx.float32)).tobytes()
                    (output/f'{key}-{part}.f32').write_bytes(data)
                    record[part+'_sha256']=hashlib.sha256(data).hexdigest()
                manifest.append(record)
                print(json.dumps(record),flush=True)
(output/'manifest.json').write_text(json.dumps(manifest))