"""Export the existing Kev vision tower directly to Core ML MIL, fixed grid.

Uses the fixed-grid/full-image-attention design in john-rocky/CoreML-LLM,
but reads our verified MLX weights and needs no PyTorch or model download.
No resampling or token reduction is performed. FP16 arithmetic differs from
the original BF16 tower and is explicitly evaluated, not assumed equivalent.
"""
import argparse,hashlib,json,time
from pathlib import Path
import numpy as np
import mlx.core as mx
from mlx.utils import tree_flatten
from mlx_vlm.models.qwen3_vl.vision import VisionModel
from mlx_vlm.models.qwen3_vl.config import VisionConfig
import coremltools as ct
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil.mil import types

ROOT=Path(__file__).resolve().parents[1];OUT=ROOT/'coreml_exports'

def main():
    p=argparse.ArgumentParser();p.add_argument('--grid',type=int,default=24)
    p.add_argument('--attention',choices=['manual','sdpa'],default='manual');p.add_argument('--out');args=p.parse_args()
    dest=OUT/(args.out or f'grid-{args.grid}-{args.attention}');dest.mkdir(parents=True,exist_ok=False)
    config=json.loads((ROOT/'config.json').read_text())['vision_config']
    assert not config['deepstack_visual_indexes']
    checkpoint=ROOT/'model.safetensors'
    h=hashlib.sha256()
    with checkpoint.open('rb') as f:
        for b in iter(lambda:f.read(8*1024**2),b''):h.update(b)
    expected=json.loads((ROOT/'provenance.json').read_text())['sha256'];assert h.hexdigest()==expected
    model=VisionModel(VisionConfig.from_dict(config))
    loaded=mx.load(str(checkpoint));model.load_weights([(k.removeprefix('vision_tower.'),v) for k,v in loaded.items() if k.startswith('vision_tower.')],strict=True)
    del loaded;model.eval();mx.eval(model.parameters());mx.clear_cache()
    weights={k:np.array(v.astype(mx.float32)).astype(np.float16) for k,v in tree_flatten(model.parameters())}
    grid=mx.array([[1,args.grid,args.grid]],dtype=mx.int32);n=args.grid**2;d=config['hidden_size'];heads=config['num_heads'];hd=d//heads
    pos=np.array(model.fast_pos_embed_interpolate(grid).astype(mx.float32)).astype(np.float16)
    freqs=model.rot_pos_emb(grid);cos=np.array(mx.tile(mx.cos(freqs),(1,2))).astype(np.float16)[:,None,:]
    sin=np.array(mx.tile(mx.sin(freqs),(1,2))).astype(np.float16)[:,None,:]
    del model;mx.clear_cache()
    (dest/'manifest.json').write_text(json.dumps({'checkpoint_sha256':expected,'grid':[1,args.grid,args.grid],
        'image_tokens':n//4,'config':config,'coremltools':ct.__version__,'attention':args.attention,
        'precision':'FP16 weights converted from original BF16; FP16 Core ML computation',
        'script_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest()},indent=2)+'\n')
    (dest/'ane_export.py').write_bytes(Path(__file__).read_bytes())
    def linear(x,name):return mb.linear(x=x,weight=weights[name+'.weight'],bias=weights[name+'.bias'],name=name.replace('.','_'))
    def norm(x,name):return mb.layer_norm(x=x,axes=[-1],gamma=weights[name+'.weight'],beta=weights[name+'.bias'],epsilon=np.float16(1e-6))
    def rope(x):
        a,b=mb.split(x=x,num_splits=2,axis=-1)
        rotated=mb.concat(values=[mb.mul(x=b,y=np.float16(-1)),a],axis=-1)
        return mb.add(x=mb.mul(x=x,y=cos),y=mb.mul(x=rotated,y=sin))
    @mb.program(input_specs=[mb.TensorSpec(shape=(n,1536),dtype=types.fp16)],opset_version=ct.target.iOS18)
    def program(pixels):
        # MLX Conv3d stores [out,time,height,width,channel]; input patches are CTHW.
        w=weights['patch_embed.proj.weight'].transpose(0,4,1,2,3).reshape(d,1536)
        x=mb.linear(x=pixels,weight=w,bias=weights['patch_embed.proj.bias'])
        x=mb.add(x=x,y=pos)
        for i in range(config['depth']):
            prefix=f'blocks.{i}';a=norm(x,prefix+'.norm1');qkv=linear(a,prefix+'.attn.qkv')
            qkv=mb.reshape(x=qkv,shape=[n,3,heads,hd]);q,k,v=mb.split(x=qkv,num_splits=3,axis=1)
            q=rope(mb.squeeze(x=q,axes=[1]));k=rope(mb.squeeze(x=k,axes=[1]));v=mb.squeeze(x=v,axes=[1])
            q,k,v=[mb.expand_dims(x=mb.transpose(x=t,perm=[1,0,2]),axes=[0]) for t in [q,k,v]]
            if args.attention=='sdpa':a=mb.scaled_dot_product_attention(query=q,key=k,value=v)
            else:
                scores=mb.matmul(x=mb.mul(x=q,y=np.float16(hd**-.5)),y=k,transpose_y=True)
                a=mb.matmul(x=mb.softmax(x=scores,axis=-1),y=v)
            a=mb.reshape(x=mb.transpose(x=a,perm=[0,2,1,3]),shape=[n,d])
            x=mb.add(x=x,y=linear(a,prefix+'.attn.proj'))
            a=norm(x,prefix+'.norm2');a=mb.gelu(x=linear(a,prefix+'.mlp.linear_fc1'),mode='TANH_APPROXIMATION')
            x=mb.add(x=x,y=linear(a,prefix+'.mlp.linear_fc2'))
        x=mb.reshape(x=norm(x,'merger.norm'),shape=[n//4,4*d])
        x=mb.gelu(x=linear(x,'merger.linear_fc1'),mode='EXACT')
        return mb.identity(x=linear(x,'merger.linear_fc2'),name='features')
    print('MIL built',flush=True);t=time.perf_counter()
    converted=ct.convert(program,convert_to='mlprogram',minimum_deployment_target=ct.target.iOS18,
        compute_precision=ct.precision.FLOAT16,compute_units=ct.ComputeUnit.CPU_AND_NE,skip_model_load=True)
    converted.save(str(dest/'vision.mlpackage'))
    print(json.dumps({'saved':str(dest),'conversion_seconds':time.perf_counter()-t}),flush=True)

if __name__=='__main__':main()