"""Fixed architecture. Imports MLX only when the head is constructed."""
def create_head(cfg):
    import mlx.core as mx
    import mlx.nn as nn
    class JointSchemaHead(nn.Module):
        def __init__(self):
            super().__init__();d=cfg['head_dims']
            self.input_norm=nn.LayerNorm(cfg['hidden_size']);self.project=nn.Linear(cfg['hidden_size'],d)
            self.evidence_query_norm=nn.LayerNorm(d);self.evidence_memory_norm=nn.LayerNorm(d)
            self.evidence_attention=nn.MultiHeadAttention(d,cfg['attention_heads'])
            self.joint=nn.TransformerEncoder(cfg['joint_layers'],d,cfg['attention_heads'],mlp_dims=cfg['ffn_dims'],dropout=0.0,activation=nn.gelu,norm_first=True)
            self.score=nn.Linear(d,1)
        def __call__(self,candidates,evidence,candidate_valid,evidence_valid):
            q=self.project(self.input_norm(candidates.astype(mx.float32)))
            m=self.project(self.input_norm(evidence.astype(mx.float32)))
            # Additive key-padding masks only. Never causal; no candidate-index embeddings.
            evidence_mask=mx.where(evidence_valid[:,None,None,:],mx.array(0.,mx.float32),mx.array(-1e9,mx.float32))
            candidate_mask=mx.where(candidate_valid[:,None,None,:],mx.array(0.,mx.float32),mx.array(-1e9,mx.float32))
            qn=self.evidence_query_norm(q);mn=self.evidence_memory_norm(m)
            q=q+self.evidence_attention(qn,mn,mn,evidence_mask)
            q=self.joint(q,mask=candidate_mask)
            logits=self.score(q).squeeze(-1)
            return mx.where(candidate_valid,logits,mx.array(-1e9,mx.float32))
    return JointSchemaHead()

def pack_numpy(features):
    import numpy as np
    if not features:raise ValueError('Empty batch')
    hidden=features[0]['candidates'].shape[-1];n=max(len(f['candidates']) for f in features);m=max(len(f['evidence']) for f in features)
    q=np.zeros((len(features),n,hidden),dtype=np.float32);e=np.zeros((len(features),m,hidden),dtype=np.float32)
    qv=np.zeros((len(features),n),dtype=bool);ev=np.zeros((len(features),m),dtype=bool);labels=[]
    for i,f in enumerate(features):
        a,b=f['candidates'],f['evidence'];label=f['label']
        if a.ndim!=2 or b.ndim!=2 or len(a)<2 or len(b)<1 or a.shape[1]!=hidden or b.shape[1]!=hidden or not 0<=label<len(a):raise ValueError('Invalid feature shape/label')
        if not np.isfinite(a).all() or not np.isfinite(b).all():raise ValueError('Nonfinite features')
        q[i,:len(a)]=a;e[i,:len(b)]=b;qv[i,:len(a)]=True;ev[i,:len(b)]=True;labels.append(label)
    return q,e,qv,ev,np.array(labels,dtype=np.int32)

def pack(features):
    import mlx.core as mx
    return tuple(mx.array(a) for a in pack_numpy(features))