"""Experimental V4.1 MLX text runtime, not an oMLX plugin.

Formula reference: DeepSeek-AI's MIT-licensed V4.1 inference/model.py,
inference/engram.py and inference/kernel.py. This implementation uses MLX
operations and optional resident weights. Vision is unsupported; native MTP
and serial greedy verification are provided separately by dspark.py.
"""
import argparse
from collections import OrderedDict
import json
import math
import os
from pathlib import Path
import struct
import time

import mlx.core as mx
import numpy as np


def rms(x, weight, eps):
    f=x.astype(mx.float32)
    return (f*mx.rsqrt(mx.mean(f*f,axis=-1,keepdims=True)+eps)*weight).astype(x.dtype)


def hc_mix(x, weight, scale, base, hc, iterations, eps, norm_eps):
    f=x.reshape(*x.shape[:-2],-1).astype(mx.float32)
    m=(f@weight.T)*mx.rsqrt(mx.mean(f*f,axis=-1,keepdims=True)+norm_eps)
    pre=mx.sigmoid(m[...,:hc]*scale[0]+base[:hc])+eps
    post=2*mx.sigmoid(m[...,hc:2*hc]*scale[1]+base[hc:2*hc])
    c=(m[...,2*hc:]*scale[2]+base[2*hc:]).reshape(*f.shape[:-1],hc,hc)
    c=mx.softmax(c,axis=-1)+eps
    c=c/(mx.sum(c,axis=-2,keepdims=True)+eps)
    for _ in range(iterations-1):
        c=c/(mx.sum(c,axis=-1,keepdims=True)+eps)
        c=c/(mx.sum(c,axis=-2,keepdims=True)+eps)
    return pre,post,c


def hc_pre(x, pre):
    return mx.sum(x.astype(mx.float32)*pre[...,None],axis=-2).astype(x.dtype)


def hc_post(x, residual, post, comb):
    return (post[...,None]*x[...,None,:]+mx.einsum('...ij,...id->...jd',comb,residual.astype(mx.float32))).astype(x.dtype)


def rotary(x, position, dim, base, original=0, factor=1, beta_fast=32, beta_slow=1, inverse=False):
    freq=1/(base**(mx.arange(0,dim,2,dtype=mx.float32)/dim))
    if original:
        def corrected(r):return dim*math.log(original/(r*2*math.pi))/(2*math.log(base))
        low=max(math.floor(corrected(beta_fast)),0);high=min(math.ceil(corrected(beta_slow)),dim-1)
        smooth=1-mx.clip((mx.arange(dim//2)-low)/max(high-low,1e-3),0,1)
        freq=freq/factor*(1-smooth)+freq*smooth
    angle=freq*position*(-1 if inverse else 1)
    tail=x[...,-dim:].astype(mx.float32).reshape(*x.shape[:-1],dim//2,2)
    a,b=tail[...,0],tail[...,1]
    out=mx.stack((a*mx.cos(angle)-b*mx.sin(angle),a*mx.sin(angle)+b*mx.cos(angle)),axis=-1)
    return mx.concatenate((x[...,:-dim],out.reshape(*x.shape[:-1],dim).astype(x.dtype)),axis=-1)


def round_fp8(x):
    a=mx.abs(x);step=2**(mx.maximum(mx.floor(mx.log2(mx.maximum(a,2**-20))),-6)-3)
    return (mx.sign(x)*mx.minimum(mx.round(a/step)*step,448)).astype(mx.float32)


def round_fp4(x):
    # Reorder even codes ahead of odd codes to implement ties-to-even.
    levels=mx.array([0,1,2,4,.5,1.5,3,6],dtype=mx.float32)
    chosen=mx.argmin(mx.abs(mx.abs(x)[...,None]-levels),axis=-1)
    return mx.sign(x)*levels[chosen]


def quantize_cache(x, bits, group, scale_format='e8m0'):
    shape=x.shape;v=x.astype(mx.float32).reshape(*shape[:-1],-1,group)
    maximum=mx.max(mx.abs(v),axis=-1,keepdims=True)
    if bits==8:
        scale=2**mx.ceil(mx.log2(mx.maximum(maximum,1e-4)/448))
        out=round_fp8(v/scale)*scale
    else:
        if scale_format=='e4m3':scale=round_fp8(mx.maximum(maximum,6*2**-9)/6)
        else:scale=2**mx.ceil(mx.log2(mx.maximum(maximum,6*2**-126)/6))
        out=round_fp4(v/scale)*scale
    return out.reshape(shape).astype(x.dtype)


def sparse_attention(q, kv, sink):
    scores=(q.astype(mx.float32)@kv.astype(mx.float32).T)*(q.shape[-1]**-.5)
    p=mx.softmax(mx.concatenate((scores,sink[:,None]),axis=-1),axis=-1)
    return (p[:,:-1]@kv.astype(mx.float32)).astype(q.dtype)


compiled_hc_mix=mx.compile(hc_mix)
compiled_rms=mx.compile(rms)


class Weights:
    """Read only requested arrays/embedding rows, bounded LRU on the GPU."""
    def __init__(self,path,cache_bytes=512*1024**2):
        self.path=Path(path);self.config=json.loads((self.path/'config.json').read_text())
        self.entries={};self.cache=OrderedDict();self.used=0;self.limit=cache_bytes
        self.resident={};self.resident_bytes=0
        self.disk_bytes_read=0;self.disk_read_calls=0;self.disk_bytes_by_key={}
        self.engram_cache_rows=0;self.engram_rows=OrderedDict()
        self.engram_cache_hits=0;self.engram_cache_misses=0
        index=json.loads((self.path/'model.safetensors.index.json').read_text())['weight_map']
        for name in sorted(set(index.values())):
            p=self.path/name
            with p.open('rb') as f:
                n=struct.unpack('<Q',f.read(8))[0];h=json.loads(f.read(n))
            for key,s in h.items():
                if key=='__metadata__':continue
                self.entries[key]=dict(s,path=p,offset=8+n+s['data_offsets'][0])
        assert set(index)==set(self.entries)
        self.q=self.config['quantization']

    @staticmethod
    def backbone_key(key):
        return '.engram.embed.' not in key and not key.startswith(('mtp.','vision.','aligner.','image_'))

    def preload_backbone(self,budget_bytes=None,progress=False):
        import psutil
        keys=[k for k in self.entries if self.backbone_key(k)]
        total=sum(self.entries[k]['data_offsets'][1]-self.entries[k]['data_offsets'][0] for k in keys)
        reserve=32*1024**3
        device=mx.device_info()
        safe=min(psutil.virtual_memory().available-reserve,
                 device['max_recommended_working_set_size']-mx.get_active_memory()-reserve)
        budget=safe if budget_bytes is None else min(safe,budget_bytes)
        if total>budget:raise MemoryError(f'Backbone needs {total} bytes; safe budget is {budget}')
        if self.resident:raise RuntimeError('Backbone already preloaded')
        start=time.monotonic();last=start
        if progress:print(json.dumps({'phase':'resident_load','bytes':total,'tensors':len(keys),'safe_budget_bytes':budget}),flush=True)
        for i,key in enumerate(keys):
            value=self.read(key,_pin=True)
            self.resident[key]=value;self.resident_bytes+=value.nbytes
            if progress and time.monotonic()-last>=15:
                print(json.dumps({'phase':'resident_load','loaded_bytes':self.resident_bytes,'completed':i+1,'total':len(keys),'seconds':time.monotonic()-start}),flush=True)
                last=time.monotonic()
        mx.clear_cache()
        if progress:print(json.dumps({'phase':'resident_ready','resident_bytes':self.resident_bytes,'seconds':time.monotonic()-start}),flush=True)

    def record_read(self,key,count):
        self.disk_bytes_read+=count;self.disk_read_calls+=1
        self.disk_bytes_by_key[key]=self.disk_bytes_by_key.get(key,0)+count

    def read(self,key,rows=None,_pin=False):
        if key in self.resident:
            value=self.resident[key]
            if rows is None:return value
            rows=np.asarray(rows,dtype=np.int64).ravel()
            assert np.all((rows>=0)&(rows<value.shape[0]))
            return value[mx.array(rows.astype(np.int32))]
        if rows is None and key in self.cache:
            self.cache.move_to_end(key);return self.cache[key]
        s=self.entries[key];dt={'U32':'<u4','U8':'u1','BF16':'<u2','F32':'<f4'}[s['dtype']]
        width=np.dtype(dt).itemsize;shape=s['shape']
        if rows is None:
            count=math.prod(shape)*width
            if count>self.limit*2 and not _pin:raise MemoryError('Whole-tensor read blocked; use row lookup for '+key)
            with s['path'].open('rb') as f:f.seek(s['offset']);blob=f.read(count)
            assert len(blob)==count
            self.record_read(key,count)
            a=np.frombuffer(blob,dtype=dt).reshape(shape).copy()
        else:
            rows=np.asarray(rows,dtype=np.int64).ravel();stride=math.prod(shape[1:])*width
            assert np.all((rows>=0)&(rows<shape[0]))
            a=np.empty((len(rows),*shape[1:]),dtype=dt)
            with s['path'].open('rb') as f:
                for i,row in enumerate(rows):
                    f.seek(s['offset']+int(row)*stride);blob=f.read(stride);assert len(blob)==stride
                    self.record_read(key,stride)
                    a[i]=np.frombuffer(blob,dtype=dt).reshape(shape[1:])
        out=mx.array(a)
        if s['dtype']=='BF16':out=out.view(mx.bfloat16)
        mx.eval(out)
        if rows is None and not _pin and out.nbytes<=self.limit:
            while self.cache and self.used+out.nbytes>self.limit:
                _,old=self.cache.popitem(last=False);self.used-=old.nbytes
            self.cache[key]=out;self.used+=out.nbytes
        return out

    def quant_config(self,base):
        cfg=self.q.get(base,self.q)
        return {k:cfg[k] for k in ('bits','group_size','mode')}

    def linear(self,base,x):
        if base+'.scales' in self.entries:
            cfg=self.quant_config(base)
            bias=self.read(base+'.biases') if base+'.biases' in self.entries else None
            y=mx.quantized_matmul(x,self.read(base+'.weight'),self.read(base+'.scales'),bias,transpose=True,**cfg)
        else:
            w=self.read(base+'.weight');y=x.astype(w.dtype)@w.T
        if base+'.bias' in self.entries:y=y+self.read(base+'.bias')
        return y.astype(x.dtype)

    def embedding(self,base,ids):
        if '.engram.embed' in base and self.engram_cache_rows>0:
            original=np.asarray(ids);flat=[int(i) for i in original.ravel()]
            missing=list(dict.fromkeys(i for i in flat if (base,i) not in self.engram_rows))
            resolved={}
            for i in flat:
                if (base,i) in self.engram_rows:
                    resolved[i]=self.engram_rows[base,i];self.engram_rows.move_to_end((base,i))
                    self.engram_cache_hits+=1
                else:self.engram_cache_misses+=1
            if missing:
                values=self._embedding_uncached(base,missing);mx.eval(values)
                for n,i in enumerate(missing):
                    resolved[i]=values[n];self.engram_rows[base,i]=values[n]
                    while len(self.engram_rows)>self.engram_cache_rows:self.engram_rows.popitem(last=False)
            result=mx.stack([resolved[i] for i in flat])
            return result.reshape(*original.shape,result.shape[-1])
        return self._embedding_uncached(base,ids)

    def _embedding_uncached(self,base,ids):
        ids=np.asarray(ids);w=self.read(base+'.weight',ids)
        if base+'.scales' in self.entries:
            s=self.read(base+'.scales',ids);b=self.read(base+'.biases',ids) if base+'.biases' in self.entries else None
            w=mx.dequantize(w,s,b,**self.quant_config(base))
        return w.reshape(*ids.shape,w.shape[-1])

    def grouped(self,base,x,groups):
        shape=self.entries[base+'.weight']['shape'];rows=shape[0]//groups
        out=[]
        for g in range(groups):
            ids=np.arange(g*rows,(g+1)*rows);w=self.read(base+'.weight',ids)
            if base+'.scales' in self.entries:
                s=self.read(base+'.scales',ids);b=self.read(base+'.biases',ids)
                y=mx.quantized_matmul(x[g:g+1],w,s,b,transpose=True,**self.quant_config(base))
            else:y=x[g:g+1]@w.T
            out.append(y)
        return mx.concatenate(out,axis=-1)


def is_prime(n):
    if n<2:return False
    if n%2==0:return n==2
    return all(n%d for d in range(3,math.isqrt(n)+1,2))


class EngramHash:
    def __init__(self,c,tokenizer):
        from tokenizers import Regex,normalizers
        normalizer=normalizers.Sequence([normalizers.NFKC(),normalizers.NFD(),normalizers.StripAccents(),normalizers.Lowercase(),normalizers.Replace(Regex(r'[ \t\r\n]+'),' '),normalizers.Replace(Regex(r'^ $'),'\ue000'),normalizers.Strip(),normalizers.Replace('\ue000',' ')])
        keys={};self.token_map=[]
        for token in range(c['vocab_size']):
            text=tokenizer.decode([token],skip_special_tokens=False)
            key=tokenizer.id_to_token(token) if '\ufffd' in text else (normalizer.normalize_str(text) or text)
            if key not in keys:keys[key]=len(keys)
            self.token_map.append(keys[key])
        assert len(keys)==c['engram_compressed_vocab_size'], 'Engram tokenizer map mismatch'
        self.pad=self.token_map[c['engram_pad_token_id']];self.history=[];self.tables={}
        seen=set();order=c['engram_max_ngram_size']
        for layer,expected in zip(c['engram_layer_ids'],c['engram_num_embeddings']):
            primes=[]
            for _ in range(order-1):
                current=c['engram_vocab_size']-1
                for _ in range(c['engram_n_heads']):
                    current+=1
                    while current in seen or not is_prime(current):current+=1
                    seen.add(current);primes.append(current)
            assert sum(primes)==expected
            rng=np.random.default_rng(10007*layer)
            bound=max(1,((2**63-1)//len(keys))//2)
            mult=rng.integers(0,bound,size=order,dtype=np.int64)*2+1
            self.tables[layer]=(primes,np.cumsum([0,*primes[:-1]]),mult)
        self.order=order;self.heads=c['engram_n_heads']

    def push(self,token):self.history.append(self.token_map[token])

    def ids(self,layer):
        primes,offsets,mult=self.tables[layer]
        tokens=[self.history[-1-i] if i<len(self.history) else self.pad for i in range(self.order)]
        rolling=int(tokens[0])*int(mult[0]);ids=[]
        for i in range(1,self.order):
            rolling^=int(tokens[i])*int(mult[i])
            for head in range(self.heads):
                j=(i-1)*self.heads+head;ids.append(rolling%primes[j]+int(offsets[j]))
        return ids


class TextRuntime:
    def __init__(self,path,max_tokens=256,resident_backbone=False,execution_mode='reference'):
        from tokenizers import Tokenizer
        self.w=Weights(path);self.c=self.w.config['text_config'];self.max_tokens=max_tokens
        assert self.w.config['model_type']=='deepseek_v41'
        self.tokenizer=Tokenizer.from_file(str(Path(path)/'tokenizer.json'))
        self.hash=EngramHash(self.c,self.tokenizer)
        self.position=0;self.states={};self.shared={}
        if execution_mode not in ('reference','deferred','compiled'):raise ValueError('Unknown execution mode')
        self.execution_mode=execution_mode
        if resident_backbone:self.w.preload_backbone(progress=True)

    def norm(self,name,x):
        fn=compiled_rms if getattr(self,'execution_mode','reference')=='compiled' else rms
        return fn(x,self.w.read(name+'.weight'),self.c['rms_norm_eps'])

    def rotate(self,x,pos,ratio,inverse=False):
        c=self.c;r=c.get('rope_scaling',{})
        return rotary(x,pos,c['qk_rope_head_dim'],c['compress_rope_theta'] if ratio else c['rope_theta'],r.get('original_max_position_embeddings',0) if ratio else 0,r.get('factor',1),r.get('beta_fast',32),r.get('beta_slow',1),inverse)

    def mixes(self,base,kind,x):
        c=self.c
        fn=compiled_hc_mix if getattr(self,'execution_mode','reference')=='compiled' else hc_mix
        return fn(x,self.w.read(base+'.hc_'+kind+'_fn'),self.w.read(base+'.hc_'+kind+'_scale'),self.w.read(base+'.hc_'+kind+'_base'),c['hc_mult'],c['hc_sinkhorn_iters'],c['hc_eps'],c['rms_norm_eps'])

    def engram(self,base,h,layer):
        ids=self.hash.ids(layer)
        values=self.w.embedding(base+'.embed',ids).reshape(1,-1)
        kv=self.w.linear(base+'.wkv',values).astype(mx.float32)
        c=self.c;dim=c['hidden_size'];hc=c['hc_mult'];eps=c['rms_norm_eps']
        key=kv[:,:hc*dim].reshape(1,hc,dim);value=kv[:,hc*dim:]
        x=h.astype(mx.float32);weight=self.w.read(base+'.q_weight').astype(mx.float32)*self.w.read(base+'.k_weight').astype(mx.float32)
        std=mx.rsqrt(mx.mean(x*x,axis=-1)+eps)*mx.rsqrt(mx.mean(key*key,axis=-1)+eps)
        dot=mx.sum(x*key*weight,axis=-1)*std*dim**-.5
        gate=mx.sigmoid(mx.where(dot<0,-1,1)*mx.sqrt(mx.maximum(mx.abs(dot),1e-6)))
        return (x+gate[...,None]*value[:,None,:]).astype(h.dtype)

    def expert(self,base,x,routing=None):
        a=self.w.linear(base+'.w1',x).astype(mx.float32);b=self.w.linear(base+'.w3',x).astype(mx.float32)
        limit=self.c['swiglu_limit']
        if limit>0:a=mx.minimum(a,limit);b=mx.clip(b,-limit,limit)
        y=a*mx.sigmoid(a)*b
        if routing is not None:y=y*routing
        return self.w.linear(base+'.w2',y.astype(x.dtype))

    def moe(self,base,x):
        c=self.c
        logits=x.astype(mx.float32)@self.w.read(base+'.gate.weight').astype(mx.float32).T
        scores=mx.sqrt(mx.logaddexp(logits,mx.zeros_like(logits)))
        k=c['num_experts_per_tok']
        picks=mx.argsort(scores+self.w.read(base+'.gate.bias'),axis=-1)[:,-k:]
        selected=mx.take_along_axis(scores,picks,axis=-1)
        if c['norm_topk_prob'] and k>1:selected=selected/(mx.sum(selected,axis=-1,keepdims=True)+1e-20)
        selected=selected*c['routed_scaling_factor'];mx.eval(picks,selected)
        result=mx.zeros_like(x).astype(mx.float32)
        for index,weight in zip(picks.tolist()[0],selected.tolist()[0]):
            result=result+self.expert(base+f'.experts.{index}',x,weight).astype(mx.float32)
            if getattr(self,'execution_mode','reference')=='reference':mx.eval(result)
        return (result+self.expert(base+'.shared_experts',x).astype(mx.float32)).astype(x.dtype)

    def attention(self,base,x,layer):
        c=self.c;t=self.position;ratio=c['compress_ratios'][layer]
        st=self.states.setdefault(layer,{'window':[],'pool':[],'score':[],'compressed':[],'index':[]})
        qr=self.norm(base+'.q_norm',self.w.linear(base+'.wq_a',x))
        q=self.w.linear(base+'.wq_b',qr).reshape(c['num_attention_heads'],c['head_dim'])
        q=self.rotate(q,t,ratio)
        kv=self.norm(base+'.kv_norm',self.w.linear(base+'.wkv',x))
        kv=quantize_cache(self.rotate(kv,t,ratio),8,32)
        st['window']=(st['window']+[kv])[-c['sliding_window']:]
        selected=mx.concatenate(st['window'],axis=0)
        latent=None
        if ratio:
            if layer in c['kv_source_layer_ids']:
                pool=self.w.linear(base+'.compressor.wkv',x.astype(mx.float32) if ratio>1 else x)
                if ratio>1:
                    st['pool'].append(pool);st['score'].append(self.w.linear(base+'.compressor.wgate',x.astype(mx.float32)))
                    if (t+1)%ratio==0:
                        a=mx.concatenate(st['pool'],axis=0);b=mx.concatenate(st['score'],axis=0)
                        latent=mx.sum(a*mx.softmax(b,axis=0),axis=0,keepdims=True).astype(x.dtype)
                        st['pool']=[];st['score']=[]
                else:latent=pool
                if latent is not None:latent=self.norm(base+'.compressor.norm',latent)
                self.shared['compressed']=st['compressed'];self.shared['index']=st['index']
            if layer in c['index_source_layer_ids']:
                if layer in c['kv_source_layer_ids'] and latent is not None:
                    ik=self.norm(base+'.indexer.k_norm',self.w.linear(base+'.indexer.wk',latent))
                    st['index'].append(quantize_cache(self.rotate(ik,t+1-ratio,ratio),4,32))
                n=(t+1)//ratio
                if n:
                    iq=self.w.linear(base+'.indexer.wq_b',qr).reshape(c['index_n_heads'],c['index_head_dim'])
                    iq=quantize_cache(self.rotate(iq,t,ratio),4,32)
                    keys=mx.concatenate(self.shared['index'],axis=0)
                    importance=self.w.linear(base+'.indexer.weights_proj',x).reshape(-1)*c['index_head_dim']**-.5*c['index_n_heads']**-.5
                    scores=mx.sum(mx.maximum(iq.astype(mx.float32)@keys.astype(mx.float32).T,0)*importance[:,None],axis=0)
                    if layer==c['candidate_source_layer_id']:
                        block=c['candidate_block_size'];padded=mx.pad(scores,(0,(-n)%block),constant_values=-float('inf'))
                        block_scores=mx.max(padded.reshape(-1,block),axis=-1)
                        block_scores=mx.where(mx.arange(len(block_scores))==(n-1)//block,float('inf'),block_scores)
                        picks=mx.argsort(block_scores)[-min(c['candidate_topk_blocks'],len(block_scores)):]
                        self.shared['candidates']=mx.any(mx.arange(n)[:,None]//block==picks[None,:],axis=-1)
                    elif layer>c['candidate_source_layer_id']:
                        scores=mx.where(self.shared['candidates'],scores,-float('inf'))
                    self.shared['topk']=mx.sort(mx.argsort(scores)[-min(c['index_topk'],n):])
                else:self.shared['topk']=mx.array([],dtype=mx.int32)
            if latent is not None:
                st['compressed'].append(quantize_cache(self.rotate(latent,t+1-ratio,ratio),4,16,'e4m3'))
            if self.shared['compressed']:
                compressed=mx.concatenate(self.shared['compressed'],axis=0)
                selected=mx.concatenate((selected,compressed[self.shared['topk']]),axis=0)
        out=sparse_attention(q,selected,self.w.read(base+'.attn_sink'))
        out=self.rotate(out,t,ratio,inverse=True).reshape(c['o_groups'],-1)
        out=self.w.grouped(base+'.wo_a',out,c['o_groups'])
        return self.w.linear(base+'.wo_b',out)

    def step(self,token,progress=False):
        if self.position>=self.max_tokens:raise ValueError('Reference context limit reached')
        if token==self.w.config.get('image_token_id'):raise NotImplementedError('Vision is not validated')
        c=self.c;self.hash.push(int(token));self.shared={}
        x=self.w.embedding('embed',[token]);h=mx.repeat(x[:,None,:],c['hc_mult'],axis=1)
        pre=mx.array([[1.]+[0.]*(c['hc_mult']-1)],dtype=mx.float32)
        targets=[]
        for layer in range(c['num_hidden_layers']):
            base=f'layers.{layer}'
            if layer in c['engram_layer_ids']:h=self.engram(base+'.engram',h,layer)
            if layer in c['dspark_target_layer_ids']:targets.append(mx.mean(h,axis=1))
            ap,post,comb=self.mixes(base,'attn',h)
            a=self.attention(base+'.attn',self.norm(base+'.attn_norm',hc_pre(h,pre)),layer)
            h=hc_post(a,h,post,comb)
            pre,post,comb=self.mixes(base,'ffn',h)
            h=hc_post(self.moe(base+'.ffn',self.norm(base+'.ffn_norm',hc_pre(h,ap))),h,post,comb)
            mx.eval(h,pre)
            if progress:print(json.dumps({'position':self.position,'layer':layer,'finite':bool(mx.all(mx.isfinite(h)))}),flush=True)
            if not bool(mx.all(mx.isfinite(h))):raise FloatingPointError('Nonfinite hidden state')
        logits=self.w.linear('head',self.norm('norm',hc_pre(h,pre)).astype(mx.float32))
        mx.eval(logits);self.position+=1
        return logits, mx.concatenate(targets,axis=-1)


def main():
    p=argparse.ArgumentParser();p.add_argument('--model',required=True);p.add_argument('--token-id',type=int);p.add_argument('--load-only',action='store_true')
    p.add_argument('--resident-backbone',action='store_true')
    p.add_argument('--execution-mode',choices=('reference','deferred','compiled'),default='compiled')
    args=p.parse_args();start=time.monotonic()
    model=TextRuntime(args.model,resident_backbone=args.resident_backbone,execution_mode=args.execution_mode)
    print(json.dumps({'adapter':'isolated-experimental-v41-text','indexed_tensors':len(model.w.entries),'mtp_tensors':sum(k.startswith('mtp.') for k in model.w.entries),'engram_map_verified':True,'omlx_integrated':False,'mtp_generation_tested':False}),flush=True)
    if not args.load_only:
        if args.token_id is None:raise ValueError('Explicit token-id required for one-token diagnostic')
        logits,_=model.step(args.token_id,progress=True)
        print(json.dumps({'next_token_id':int(mx.argmax(logits)),'seconds':time.monotonic()-start,'peak_MLX_bytes':mx.get_peak_memory(),'quality_validated':False}),flush=True)


if __name__=='__main__':main()