"""Lossless separate-expert layout adapter for pinned official Qwen3.5 MoE.

Usage: with native_layout_adapter() as audit: mlx_lm.convert.convert(...)
The context changes only source layout handling and restores the original method.
All quantization, model math, serialization and reload remain official MLX code.
"""
import contextlib,importlib.metadata,json,pathlib,re,struct,threading
_PATTERN=re.compile(r'^(model\.language_model|language_model\.model)\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$')
_LOCK=threading.Lock()

def _plan(weights,layers,experts,hidden,intermediate):
    found={};recognized=set()
    for key,value in weights.items():
        match=_PATTERN.fullmatch(key)
        if match:
            base,layer,expert,proj=match.groups();layer=int(layer);expert=int(expert)
            assert key==f'{base}.layers.{layer}.mlp.experts.{expert}.{proj}.weight','Noncanonical expert index'
            assert 0<=layer<layers and 0<=expert<experts,('Out-of-range expert key',key)
            shape=tuple(value.shape) if hasattr(value,'shape') else tuple(value['shape'])
            dtype=str(value.dtype) if hasattr(value,'dtype') else value['dtype']
            required=(hidden,intermediate) if proj=='down_proj' else (intermediate,hidden)
            assert shape==required,('Expert shape mismatch',key,shape,required)
            found.setdefault((base,layer,proj),{})[expert]=(key,shape,dtype);recognized.add(key)
        elif re.match(r'^(model\.language_model|language_model\.model)\.layers\.\d+\.mlp\.experts\.',key) and re.search(r'\.experts\.\d+\.',key):
            raise ValueError('Unexpected independent expert key: '+key)
    if not found:return []
    bases={g[0] for g in found};assert len(bases)==1,'Mixed expert key roots';base=next(iter(bases))
    assert len(found)==layers*3,'Missing layer or projection group'
    plan=[]
    for layer in range(layers):
        prefix=f'{base}.layers.{layer}.mlp'
        assert not any(k.startswith(prefix+'.switch_mlp.') or k in (prefix+'.experts.gate_up_proj',prefix+'.experts.down_proj') for k in weights),'Mixed packed/independent expert layouts'
        for proj in ('gate_proj','up_proj','down_proj'):
            group=found[(base,layer,proj)];assert set(group)==set(range(experts)),('Missing/extra experts',layer,proj)
            assert len({v[2] for v in group.values()})==1,('Mixed expert dtypes',layer,proj)
            plan.append((prefix+'.switch_mlp.'+proj+'.weight',[group[i][0] for i in range(experts)]))
    return plan

def stack_independent_experts(weights,args):
    import mlx.core as mx
    plan=_plan(weights,args.num_hidden_layers,args.num_experts,args.hidden_size,args.moe_intermediate_size)
    if not plan:return weights,{'adapted':False,'groups':0,'source_expert_tensors':0}
    result=dict(weights)
    for destination,keys in plan:
        # Numeric expert order is explicit; mx.stack does not cast or quantize.
        result[destination]=mx.stack([result.pop(k) for k in keys],axis=0)
    assert len(result)==len(weights)-sum(len(keys)-1 for _,keys in plan)
    return result,{'adapted':True,'groups':len(plan),'source_expert_tensors':sum(len(k) for _,k in plan),'expert_order':'numeric 0..num_experts-1','value_transform':'stack only; no cast, quantization, norm or conv transform'}

@contextlib.contextmanager
def native_layout_adapter():
    assert importlib.metadata.version('mlx-lm')=='0.31.3','Adapter audited for mlx-lm 0.31.3 only'
    from mlx_lm.models.qwen3_5_moe import Model
    if not _LOCK.acquire(blocking=False):raise RuntimeError('Layout adapter already active in this process')
    original=Model.sanitize;audit={'calls':[],'native_sanitize_calls':0,'restored':False}
    def sanitize(self,weights):
        transformed,row=stack_independent_experts(weights,self.language_model.args)
        audit['calls'].append(row);audit['native_sanitize_calls']+=1
        return original(self,transformed)
    try:Model.sanitize=sanitize;yield audit
    finally:Model.sanitize=original;audit['restored']=True;_LOCK.release()

def source_header_preflight(source):
    source=pathlib.Path(source);config=json.loads((source/'config.json').read_text())['text_config'];index=json.loads((source/'model.safetensors.index.json').read_text())['weight_map'];weights={}
    for name in sorted(set(index.values())):
        assert pathlib.Path(name).name==name
        with (source/name).open('rb') as f:
            size=struct.unpack('<Q',f.read(8))[0];assert size<100_000_000;header=json.loads(f.read(size))
        for key,value in header.items():
            if key=='__metadata__':continue
            assert key not in weights and index.get(key)==name;weights[key]=value
    assert set(weights)==set(index)
    plan=_plan(weights,config['num_hidden_layers'],config['num_experts'],config['hidden_size'],config['moe_intermediate_size'])
    assert len(plan)==120 and sum(len(keys) for _,keys in plan)==30720
    return {'passed':True,'source_index_tensors':len(weights),'layers':config['num_hidden_layers'],'experts_per_layer':config['num_experts'],'groups':len(plan),'independent_expert_tensors':30720,'all_shapes_checked':True,'each_group_dtype_uniform':True,'weight_payload_read':False,'model_instantiated':False}