"""Attention functions and discrete diagnostics, no continuous alignment/fitting."""
import numpy as np
from functional.model import StreamingLM,PREFIX,attention,rope,sigmoid


def assignment(cost):
    """Exact minimum-cost bijection via O(n^2 2^n) dynamic programming."""
    cost=np.asarray(cost,dtype=np.float64);n=cost.shape[0]
    assert cost.shape==(n,n) and n<=12 and np.isfinite(cost).all()
    states={0:(0.,())}
    for i in range(n):
        nxt={}
        for mask,(value,perm) in states.items():
            for j in range(n):
                if mask&(1<<j):continue
                key=mask|(1<<j);candidate=(value+float(cost[i,j]),perm+(j,))
                if key not in nxt or candidate<nxt[key]:nxt[key]=candidate
        states=nxt
    return list(states[(1<<n)-1][1])


def inverse(perm):
    out=[0]*len(perm)
    for i,j in enumerate(perm):out[j]=i
    return out


def output_metrics(a,b):
    """Per-text directional error, cosine and output RMS. No fitted scaling."""
    x=a.reshape(a.shape[0],-1).astype(np.float64)
    y=b.reshape(b.shape[0],-1).astype(np.float64)
    xx=np.sum(x*x,-1);yy=np.sum(y*y,-1);diff=np.sum((x-y)**2,-1)
    return {'error_a':np.sqrt(diff/np.maximum(xx,1e-30)),
            'error_b':np.sqrt(diff/np.maximum(yy,1e-30)),
            'cosine':np.sum(x*y,-1)/np.maximum(np.sqrt(xx*yy),1e-30),
            'rms_a':np.sqrt(xx/x.shape[1]),'rms_b':np.sqrt(yy/y.shape[1])}


def distribution_metrics(a,b):
    """Aligned heads; primary queries have at least 8 legal keys (t>=7)."""
    a,b=a.astype(np.float64),b.astype(np.float64)
    assert a.shape==b.shape and a.shape[-1]>=8
    m=(a+b)*.5
    def term(x):return np.where(x>0,x*np.log(np.maximum(x,1e-300)/np.maximum(m,1e-300)),0)
    js=np.maximum(.5*np.sum(term(a)+term(b),axis=-1),0)
    eligible=js[...,7:]
    overlaps=[]
    for t in range(7,a.shape[-2]):
        aa=np.argsort(a[:,:,t,:t+1],axis=-1)[...,-5:]
        bb=np.argsort(b[:,:,t,:t+1],axis=-1)[...,-5:]
        overlaps.append(np.mean((aa[...,None]==bb[...,None,:]).any(-1),-1))
    # First-key ablation is diagnostic, not an unreported filter of positional repetitions.
    ar,br=a[...,1:],b[...,1:];ma,mb=ar.sum(-1),br.sum(-1)
    valid=(ma>1e-8)&(mb>1e-8)
    ar=ar/np.maximum(ma[...,None],1e-300);br=br/np.maximum(mb[...,None],1e-300)
    mr=(ar+br)*.5
    cond=.5*np.sum(np.where(ar>0,ar*np.log(np.maximum(ar,1e-300)/np.maximum(mr,1e-300)),0)+np.where(br>0,br*np.log(np.maximum(br,1e-300)/np.maximum(mr,1e-300)),0),-1)
    sink=[]
    for bi in range(a.shape[0]):
        which=valid[bi,:,7:];values=cond[bi,:,7:]
        sink.append(float(np.maximum(values[which],0).mean()) if which.any() else None)
    diag=np.stack([a[:,:,t,t] for t in range(7,a.shape[-2])],-1).mean((1,2))
    local=np.stack([a[:,:,t,max(0,t-2):t+1].sum(-1) for t in range(7,a.shape[-2])],-1).mean((1,2))
    return {'js':eligible.mean((1,2)),'js_p95':np.quantile(eligible.reshape(a.shape[0],-1),.95,axis=1),
            'top5':np.stack(overlaps,-1).mean((1,2)),'js_all_queries':js.mean((1,2)),
            'first_key_mass_a':a[:,:,7:,0].mean((1,2)),'first_key_mass_b':b[:,:,7:,0].mean((1,2)),
            'self_mass_a':diag,'recent3_mass_a':local,'js_without_first_key':sink}


def cross_head_output_metrics(a,b):
    """All head pairs, per text. a/b: [B,T,H,Dhidden]."""
    aa=a.transpose(0,2,1,3).reshape(a.shape[0],a.shape[2],-1).astype(np.float64)
    bb=b.transpose(0,2,1,3).reshape(b.shape[0],b.shape[2],-1).astype(np.float64)
    xx=np.sum(aa*aa,-1);yy=np.sum(bb*bb,-1)
    dot=np.einsum('bif,bjf->bij',aa,bb,optimize=True)
    diff=np.maximum(xx[:,:,None]+yy[:,None,:]-2*dot,0)
    return {'error_a':np.sqrt(diff/np.maximum(xx[:,:,None],1e-30)),
            'error_b':np.sqrt(diff/np.maximum(yy[:,None,:],1e-30)),
            'cosine':dot/np.maximum(np.sqrt(xx[:,:,None]*yy[:,None,:]),1e-30),
            'rms_a':np.broadcast_to(np.sqrt(xx/aa.shape[-1])[:,:,None],dot.shape),
            'rms_b':np.broadcast_to(np.sqrt(yy/bb.shape[-1])[:,None,:],dot.shape)}


def value_gram_cka(a,b):
    """Centered token Gram agreement, diagnostic only; not interchangeability."""
    a=a.astype(np.float64)-a.mean(axis=1,keepdims=True,dtype=np.float64)
    b=b.astype(np.float64)-b.mean(axis=1,keepdims=True,dtype=np.float64)
    ga=np.einsum('btd,bsd->bts',a,a);gb=np.einsum('btd,bsd->bts',b,b)
    return np.sum(ga*gb,(1,2))/np.maximum(np.sqrt(np.sum(ga*ga,(1,2))*np.sum(gb*gb,(1,2))),1e-30)


class QKVModel(StreamingLM):
    def project(self,context,gate,layer):
        return self.linear((context*sigmoid(gate)).reshape(*context.shape[:2],-1),f'layers.{layer}.self_attn.o_proj')

    def project_heads(self,context,gate,layer):
        """Each head's contribution in the common residual output coordinates."""
        c=context*sigmoid(gate);b,t,h,d=c.shape
        weight=self.tensors[PREFIX+f'layers.{layer}.self_attn.o_proj.weight']
        outdim,indim=weight.shape;assert indim==h*d
        out=np.empty((b,t,h,outdim),dtype=np.float32)
        for start in range(0,outdim,256):
            self.guard();end=min(start+256,outdim)
            w=weight.read(slice(start*indim,end*indim)).reshape(end-start,indim)
            for head in range(h):out[:,:,head,start:end]=c[:,:,head,:]@w[:,head*d:(head+1)*d].T
            del w;weight.evict()
        return out

    def parts(self,h,layer,heads=True):
        p=f'layers.{layer}.self_attn.';c=self.config;d=c['head_dim'];nh=c['num_attention_heads'];nk=c['num_key_value_heads']
        mixed=self.linear(h,p+'q_proj').reshape(*h.shape[:2],nh,2*d)
        q,gate=mixed[...,:d],mixed[...,d:].copy()
        k=self.linear(h,p+'k_proj').reshape(*h.shape[:2],nk,d)
        v=self.linear(h,p+'v_proj').reshape(*h.shape[:2],nk,d)
        q=self.normalization(q,p+'q_norm');k=self.normalization(k,p+'k_norm')
        r=c['rope_parameters'];q,k=(rope(x,r['rope_theta'],r['partial_rotary_factor']) for x in (q,k))
        context,probs=attention(q,k,v)
        out={'q':q,'k':k,'p':probs,'v':v,'gate':gate,'context':context,'output':self.project(context,gate,layer)}
        if heads:out['heads']=self.project_heads(context,gate,layer)
        return out

    def transplant(self,own,donor,layer,kind,target_head=None,donor_head=None,perm=None):
        repeats=self.config['num_attention_heads']//self.config['num_key_value_heads']
        if kind=='full_attention':return donor['output']
        if kind=='qkv_head':return own['output']-own['heads'][:,:,target_head]+donor['heads'][:,:,donor_head]
        probs=own['p'].copy();v=own['v'].copy()
        if kind=='qk_map':probs=donor['p'][:,perm]
        elif kind=='qk_head':probs[:,target_head]=donor['p'][:,donor_head]
        elif kind=='v_group':v[:,:,target_head]=donor['v'][:,:,donor_head]
        else:raise ValueError(kind)
        context=np.einsum('bhts,bshd->bthd',probs,np.repeat(v,repeats,axis=2))
        return self.project(context,own['gate'],layer)

    def head_transplant_output(self,own,donor,layer,target_head,donor_head):
        """QK-only transplanted target head contribution; all other heads omitted."""
        repeats=self.config['num_attention_heads']//self.config['num_key_value_heads']
        c=own['context'].copy()
        c[:,:,target_head]=np.einsum('bts,bsd->btd',donor['p'][:,donor_head],own['v'][:,:,target_head//repeats])
        # Zero remaining head contexts before projection to isolate this learned subfunction.
        for head in range(c.shape[2]):
            if head!=target_head:c[:,:,head]=0
        return self.project(c,own['gate'],layer)

    def value_group_output(self,own,donor,layer,target_group,donor_group):
        repeats=self.config['num_attention_heads']//self.config['num_key_value_heads'];c=np.zeros_like(own['context'])
        for head in range(target_group*repeats,(target_group+1)*repeats):
            c[:,:,head]=np.einsum('bts,bsd->btd',own['p'][:,head],donor['v'][:,:,donor_group])
        return self.project(c,own['gate'],layer)


class InterventionLM(QKVModel):
    def configure(self,candidate,reverse=False,cached=False):
        self.candidate=candidate;self.reverse=reverse;self.use_cache=cached;self.donor_cache=None
        self.target_layer=candidate['j'] if reverse else candidate['i']
        self.donor_layer=candidate['i'] if reverse else candidate['j']
        self.target_head=candidate.get('b') if reverse else candidate.get('a')
        self.donor_head=candidate.get('a') if reverse else candidate.get('b')
        self.perm=inverse(candidate['perm']) if reverse and candidate.get('perm') else candidate.get('perm')
        if cached:assert self.donor_layer<self.target_layer,'Cannot reuse future layer state'

    def full_attention(self,h,layer,shared_p=None):
        if not hasattr(self,'candidate'):return super().full_attention(h,layer,shared_p)
        if self.use_cache and layer==self.donor_layer:
            self.donor_cache=self.parts(h,layer,heads=self.candidate['kind']=='qkv_head')
            return self.donor_cache['output'],self.donor_cache['p']
        if layer!=self.target_layer:return super().full_attention(h,layer,shared_p)
        own=self.parts(h,layer,heads=self.candidate['kind']=='qkv_head')
        donor=self.donor_cache if self.use_cache else self.parts(h,self.donor_layer,heads=self.candidate['kind']=='qkv_head')
        assert donor is not None
        out=self.transplant(own,donor,layer,self.candidate['kind'],self.target_head,self.donor_head,self.perm)
        return out,own['p']