"""Heterogeneous LFM3 decoder stack with exact source-operation transport.

Machine cells use the compiled circuit's original sparse ReLU/head operations.
No SiLU approximation, normalization, widening or numerical conversion is added.
Language cache indices remain those of the unmodified language layers.
"""
from collections import OrderedDict
import inspect,weakref
import torch
from torch import nn
from .computer.circuit import CompiledComputer,CircuitStage

class MachineDecoderLayer(nn.Module):
    decoder_role='machine'
    def language_identity(self,hidden_states):return hidden_states

class ComputerControlDecoderLayer(MachineDecoderLayer):
    decoder_role='computer_control'
    def __init__(self,computer):
        super().__init__();self.spec=computer.spec
        self.embedding=computer.embedding;self.output=computer.output;self.programs=computer.programs
    def forward(self,hidden_states=None,*,phase='language',args=(),kwargs=None):
        if phase=='language':return hidden_states
        port=self._port();kwargs=kwargs or {}
        if phase=='boot':return CompiledComputer.boot(port,*args,**kwargs)
        if phase=='embed':return CompiledComputer.embed(port,*args,**kwargs)
        raise ValueError('Unknown control phase: '+phase)

class OrderedSparseDownProjection(nn.Module):
    """Exact w2 adapter: retain coefficient multiplication and scatter order."""
    def __init__(self,data):super().__init__();self.data=data
    def forward(self,values,template):
        values=values*self.data.coefficient
        return torch.zeros_like(template).index_add_(1,self.data.output,values)

class ExactGatedMLP(nn.Module):
    """LFM gated-MLP decomposition with exact ReLU/sparse projection adapters."""
    def __init__(self,operation):
        super().__init__();self.w1=operation.b;self.w3=operation.a
        self.w2=OrderedSparseDownProjection(operation.data);self.activation=nn.ReLU()
    def forward(self,hidden_states,active):
        values=self.w3(hidden_states)*self.activation(self.w1(hidden_states))
        return self.w2(values,hidden_states)*active[:,None]

class ExactRecurrentAttention(nn.Module):
    """Exact memory attention adapter; persistent writes are source operations."""
    def __init__(self,operation):
        super().__init__();self.spec=operation.spec;self.data=operation.data
        for name in operation.spec['maps']:self.add_module(name,getattr(operation,name))
    def forward(self,x,state,active):
        d=self.data;memory,history=state;flat=torch.cat((x,history),dim=1)
        B,V=x.shape[0],d.khead.numel()
        q=self.q(flat).relu().reshape(B,-1,128)[:,d.khead]
        k=self.k(flat).relu().reshape(B,-1,128)[:,d.khead]
        v=self.v(flat).relu().reshape(B,V,128)
        z=self.z(flat).relu().reshape(B,V,128)
        a,beta=self.a(flat),self.b(flat)+d.bconst
        factor=torch.where(a==0,1.,torch.where(a==110,.5,0.))
        updated=memory*factor[:,:,None,None];row=k.argmax(-1)
        write=(k.amax(-1)>0)&(beta==1)
        selector=torch.arange(128,device=x.device)[None,None,:]==row[:,:,None]
        updated=torch.where((write[:,:,None]&selector)[:,:,:,None],v[:,:,None,:],updated)
        memory.copy_(torch.where(active[:,None,None,None],updated,memory))
        read=((q>0)[:,:,:,None]*memory).sum(2)
        reference=read[:,:,:2].square().sum(-1,keepdim=True)
        denominator=(reference/128.).sqrt()
        normalized=torch.where(reference>0,read/torch.where(denominator>0,denominator,1.),0.)
        delta=(normalized*z*d.oweight).reshape(B,-1)
        return torch.zeros_like(x).index_add_(1,d.output.flatten(),delta)*active[:,None]
    def commit_history(self,x,state,active):
        memory,history=state
        new_history=torch.cat((x,history[:,:2*x.shape[1]]),dim=1)
        history.copy_(torch.where(active[:,None],new_history,history))

class ComputerExecutionDecoderLayer(MachineDecoderLayer):
    decoder_role='computer_execution'
    def __init__(self,operations):
        super().__init__();self.stage_specs=[operation.spec for operation in operations];self.triton_widths={}
        self.operator_norm=nn.Identity();self.ffn_norm=nn.Identity();self.use_sandwich_norm=False
        head=next((op for op in operations if op.kind=='head'),None);ffn=next((op for op in operations if op.kind=='ffn'),None)
        self.self_attn=ExactRecurrentAttention(head) if head is not None else None
        self.feed_forward=ExactGatedMLP(ffn) if ffn is not None else None
    def forward(self,hidden_states,state=None,active=None,backend='torch',scratch=None,*,phase='language'):
        if active is None:return hidden_states
        if backend=='triton':
            for spec in self.stage_specs:
                view=StageAdapter(self,spec)
                if view.kind not in self.triton_widths:raise RuntimeError("Call circuit.prepare('triton') before CUDA capture.")
                view.widths=self.triton_widths[view.kind]
                hidden_states,_=CircuitStage.triton(view,hidden_states,state if view.kind=='head' else None,active,scratch)
            return hidden_states,state
        # Same operator/residual then gated-MLP/residual organization as LFM3.
        # Missing machine sublayers are pruned, not simulated with add-zero.
        if self.self_attn is not None:
            residual=hidden_states
            attention_output=self.self_attn(self.operator_norm(hidden_states),state,active)
            hidden_states=residual+attention_output
            self.self_attn.commit_history(residual,state,active)
        if self.feed_forward is not None:
            residual=hidden_states
            ffn_output=self.feed_forward(self.ffn_norm(hidden_states),active)
            hidden_states=residual+ffn_output
        return hidden_states,state

class StageAdapter:
    """Parameter-free view for stable logical state names and program boot."""
    def __init__(self,cell,spec):self.cell=cell;self.spec=spec;self.kind=spec['kind'];self.widths=None
    @property
    def data(self):return self.cell.feed_forward.w2.data if self.kind=='ffn' else self.cell.self_attn.data
    def __getattr__(self,name):
        if self.kind=='ffn' and name in ('a','b'):return getattr(self.cell.feed_forward,'w3' if name=='a' else 'w1')
        if self.kind=='head':return getattr(self.cell.self_attn,name)
        raise AttributeError(name)
    def logical_modules(self):
        for name in self.spec['maps']:yield name,getattr(self,name)
        yield 'data',self.data
    def prepare_triton(self):
        CircuitStage.prepare_triton(self);self.cell.triton_widths[self.kind]=self.widths

class MemoryDecoderLayer(MachineDecoderLayer):
    decoder_role='memory'
    def __init__(self,ram,files):super().__init__();self.ram=ram;self.files=files
    def forward(self,hidden_states=None,*,phase='language',component='ram',method='forward',args=(),kwargs=None):
        if phase=='language':return hidden_states
        module=getattr(self,component)
        return module(*args,**(kwargs or {})) if method=='forward' else getattr(module,method)(*args,**(kwargs or {}))

class ServiceDecoderLayer(MachineDecoderLayer):
    def __init__(self,module,role):super().__init__();self.payload=module;self.decoder_role=role
    def forward(self,hidden_states=None,*,phase='language',method='forward',args=(),kwargs=None):
        if phase=='language':return hidden_states
        return self.payload(*args,**(kwargs or {})) if method=='forward' else getattr(self.payload,method)(*args,**(kwargs or {}))

class DecoderModelMixin:
    def forward(self,*args,execution_phase=None,**kwargs):
        if execution_phase is None or execution_phase=='language':return super().forward(*args,**kwargs)
        if args:raise ValueError('Machine execution uses explicit typed arguments')
        if execution_phase=='computer':
            state=kwargs.pop('computer_state')
            if kwargs:raise ValueError('Unexpected computer execution fields')
            # The original transition function calls the actual decoder cells
            # through the stage view. No separate computer module is retained.
            port=self._computer_port
            active=(state.current!=port.spec['halt'])&state.valid
            x=port.embed(state.current);scratch=torch.zeros_like(x) if port.backend=='triton' else None
            hi=0
            for cell in self.layers[1:1+self.computer_block_count]:
                head=state.heads[hi] if cell.self_attn is not None else None
                x,_=cell(x,head,active,port.backend,scratch)
                hi+=cell.self_attn is not None
            out=port.output
            family=(x[:,out.family].round()*out.family_weight).sum(1).long()
            payload=(x[:,out.payload].round()*out.payload_weight).sum(1).long()
            f=family.clamp(0,port.embedding.start.numel()-1)
            valid=((family>=0)&(family<port.embedding.start.numel())&(payload>=0)&(payload<port.embedding.size[f]))
            state.valid.copy_(state.valid&(valid|~active));predicted=port.embedding.start[f]+payload
            state.current.copy_(torch.where(active,predicted,state.current));return state.current
        index=kwargs.pop('decoder_index');method=kwargs.pop('decoder_method','forward')
        call_args=kwargs.pop('decoder_args',());call_kwargs=kwargs.pop('decoder_kwargs',{})
        component=kwargs.pop('decoder_component',None)
        if kwargs:raise ValueError('Unexpected decoder execution fields')
        layer=self.layers[index]
        if isinstance(layer,ComputerControlDecoderLayer):return layer(phase=execution_phase,args=call_args,kwargs=call_kwargs)
        if isinstance(layer,MemoryDecoderLayer):return layer(phase=execution_phase,component=component,method=method,args=call_args,kwargs=call_kwargs)
        return layer(phase=execution_phase,method=method,args=call_args,kwargs=call_kwargs)

class StageView:
    def __init__(self,owner,specs,mapping):self.owner=weakref.ref(owner);self.specs=specs;self.mapping=mapping
    def __len__(self):return len(self.specs)
    def __iter__(self):return iter(tuple(self[i] for i in range(len(self))))
    def __getitem__(self,index):
        if isinstance(index,slice):return list(self)[index]
        return StageAdapter(self.owner().layers[self.mapping[index]],self.specs[index])

class ComputerPort:
    """Parameter-free compatibility API; execution always enters decoder.forward."""
    def __init__(self,owner,spec,mapping,backend):
        self.owner=weakref.ref(owner);self.spec=spec;self.backend=backend;self.stages=StageView(owner,spec['stages'],mapping)
    @property
    def embedding(self):return self.owner().layers[0].embedding
    @property
    def output(self):return self.owner().layers[0].output
    @property
    def programs(self):return self.owner().layers[0].programs
    @property
    def device(self):return self.embedding.base.device
    def __call__(self,state):return self.owner()(execution_phase='computer',computer_state=state)
    def embed(self,tokens):return self.owner()(execution_phase='embed',decoder_index=0,decoder_args=(tokens,))
    def boot(self,*args,**kwargs):return self.owner()(execution_phase='boot',decoder_index=0,decoder_args=args,decoder_kwargs=kwargs)
    def prepare(self,backend='torch'):return CompiledComputer.prepare(self,backend)
    def _logical_modules(self):
        yield 'embedding',self.embedding;yield 'output',self.output
        for i,stage in enumerate(self.stages):
            for name,module in stage.logical_modules():yield 'stages.'+str(i)+'.'+name,module
        yield 'programs',self.programs
    def state_dict(self,destination=None,prefix='',keep_vars=False):
        result=OrderedDict() if destination is None else destination
        for name,module in self._logical_modules():module.state_dict(destination=result,prefix=prefix+name+'.',keep_vars=keep_vars)
        return result
    def named_parameters(self,prefix='',recurse=True,remove_duplicate=True):
        seen=set()
        for name,module in self._logical_modules():
            for key,value in module.named_parameters(recurse=recurse,remove_duplicate=remove_duplicate):
                if remove_duplicate and id(value) in seen:continue
                seen.add(id(value));yield ((prefix+'.') if prefix else '')+name+'.'+key,value
    def parameters(self,recurse=True):return (value for _,value in self.named_parameters(recurse=recurse))

class ComponentPort:
    def __init__(self,owner,index,component=None):self.owner=weakref.ref(owner);self.index=index;self.component=component
    def _module(self):
        layer=self.owner().layers[self.index]
        return getattr(layer,self.component) if self.component else layer.payload
    def _call(self,method,args,kwargs):
        return self.owner()(execution_phase='service',decoder_index=self.index,decoder_component=self.component,decoder_method=method,decoder_args=args,decoder_kwargs=kwargs)
    def __call__(self,*args,**kwargs):return self._call('forward',args,kwargs)
    def __getattr__(self,name):
        value=getattr(self._module(),name)
        if callable(value) and not isinstance(value,nn.Module) and name not in {'state_dict','named_parameters','parameters','named_buffers','buffers','named_modules','modules','children','named_children','get_submodule'}:
            return lambda *args,**kwargs:self._call(name,args,kwargs)
        return value

def install_decoder_stack(model):
    backbone=model.language_model.model
    if getattr(backbone,'computer_decoder_count',0):raise ValueError('Decoder stack already installed')
    original=model.computer;stage_count=len(original.stages)
    groups=[];i=0
    while i<stage_count:
        group=[i]
        if original.stages[i].kind=='head' and i+1<stage_count and original.stages[i+1].kind=='ffn':group.append(i+1)
        groups.append(group);i+=len(group)
    mapping=[0]*stage_count
    for index,group in enumerate(groups,1):
        for source in group:mapping[source]=index
    prefix=[ComputerControlDecoderLayer(original)]+[ComputerExecutionDecoderLayer([original.stages[j] for j in group]) for group in groups]
    indices={'memory':len(prefix)};prefix.append(MemoryDecoderLayer(model.ram,model.files))
    if model.ram_bridge is not None:indices['bridge']=len(prefix);prefix.append(ServiceDecoderLayer(model.ram_bridge,'bridge'))
    indices['analytics']=len(prefix);prefix.append(ServiceDecoderLayer(model.analytics,'analytics'))
    language_layers=list(backbone.layers);backbone.layers=nn.ModuleList(prefix+language_layers)
    backbone.computer_decoder_count=len(prefix);backbone.computer_block_count=len(groups);backbone.language_decoder_count=len(language_layers)
    backbone.__class__=type('LFMputerDecoderModel',(DecoderModelMixin,type(backbone)),{})
    computer=ComputerPort(backbone,original.spec,mapping,original.backend)
    object.__setattr__(backbone,'_computer_port',computer);object.__setattr__(prefix[0],'_port',weakref.ref(computer))
    for name in ['computer','ram','files','analytics']:
        model._modules.pop(name)
    if model.ram_bridge is not None:model._modules.pop('ram_bridge')
    object.__setattr__(model,'computer',computer)
    object.__setattr__(model,'ram',ComponentPort(backbone,indices['memory'],'ram'))
    object.__setattr__(model,'files',ComponentPort(backbone,indices['memory'],'files'))
    object.__setattr__(model,'analytics',ComponentPort(backbone,indices['analytics']))
    if 'bridge' in indices:object.__setattr__(model,'ram_bridge',ComponentPort(backbone,indices['bridge']))
    model.decoder_layout={'computer_stage_count':stage_count,'computer_block_count':len(groups),'stage_groups':groups,'stage_to_decoder':mapping,'computer_stage_kinds':[s.kind for s in original.stages],'machine_layer_count':len(prefix),'language_layer_count':len(language_layers),'total_layer_count':len(backbone.layers),'memory_index':indices['memory'],'bridge_index':indices.get('bridge'),'analytics_index':indices['analytics'],'roles':[layer.decoder_role for layer in prefix]+['language']*len(language_layers),'language_cache_indices':'unchanged; machine state is explicit'}
    model.decoder_layout.update(machine_hidden_size=original.spec['dims'],language_hidden_size=backbone.config.hidden_size)
    if hasattr(model,'lora_names'):
        model.lora_names=[logical_to_decoder_key('language_model.'+name,model.decoder_layout)[len('language_model.'):] for name in model.lora_names]
    model._runners={};return model

def logical_to_decoder_key(key,layout):
    base='language_model.model.layers.';p=layout['machine_layer_count']
    if key.startswith(base):
        i,tail=key[len(base):].split('.',1);return base+str(int(i)+p)+'.'+tail
    if key.startswith('computer.stages.'):
        i,tail=key[len('computer.stages.'):].split('.',1);kind=layout['computer_stage_kinds'][int(i)]
        if kind=='ffn':
            name,rest=tail.split('.',1);tail={'a':'feed_forward.w3','b':'feed_forward.w1','data':'feed_forward.w2.data'}[name]+'.'+rest
        else:tail='self_attn.'+tail
        return base+str(layout['stage_to_decoder'][int(i)])+'.'+tail
    for name in ['embedding','output','programs']:
        prefix='computer.'+name+'.'
        if key.startswith(prefix):return base+'0.'+name+'.'+key[len(prefix):]
    for prefix,index,child in [('ram.',layout['memory_index'],'ram.'),('files.',layout['memory_index'],'files.'),('ram_bridge.',layout.get('bridge_index'),'payload.'),('analytics.',layout['analytics_index'],'payload.')]:
        if key.startswith(prefix):
            if index is None:raise ValueError('Bridge tensor without a bridge layer')
            return base+str(index)+'.'+child+key[len(prefix):]
    return key

def decoder_to_logical_key(key,layout):
    base='language_model.model.layers.'
    if not key.startswith(base):return key
    i,tail=key[len(base):].split('.',1);i=int(i);p=layout['machine_layer_count']
    if i>=p:return base+str(i-p)+'.'+tail
    if i==0:return 'computer.'+tail
    if 1<=i<=layout['computer_block_count']:
        kind='ffn' if tail.startswith('feed_forward.') else 'head'
        source=next(j for j,n in enumerate(layout['stage_to_decoder']) if n==i and layout['computer_stage_kinds'][j]==kind)
        if kind=='ffn':
            for before,after in [('feed_forward.w3.','a.'),('feed_forward.w1.','b.'),('feed_forward.w2.data.','data.')]:
                if tail.startswith(before):return 'computer.stages.'+str(source)+'.'+after+tail[len(before):]
            raise ValueError('Unmapped machine FFN tensor: '+key)
        assert tail.startswith('self_attn.');return 'computer.stages.'+str(source)+'.'+tail[len('self_attn.'):]
    for index,child,prefix in [(layout['memory_index'],'ram.','ram.'),(layout['memory_index'],'files.','files.'),(layout.get('bridge_index'),'payload.','ram_bridge.'),(layout['analytics_index'],'payload.','analytics.')]:
        if i==index and tail.startswith(child):return prefix+tail[len(child):]
    raise ValueError('Unmapped decoder tensor: '+key)